Image_Inpaint_CelebHQ_Diffusions.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4",
      "authorship_tag": "ABX9TyNe9GVJ+Of1gETqd/TR22DW",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "e9a7aeae5a474eab8f60ee3046d7cebc": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_c3ec9930f8cd43f080e0a0da53a97faf",
              "IPY_MODEL_b7c5f4ed54bb4cfa98d402302891ac4f",
              "IPY_MODEL_fa98e160e55146ba86ceb9aa137e07f7"
            ],
            "layout": "IPY_MODEL_2bcea0ab312b4ea8a474bebea04edeae"
          }
        },
        "c3ec9930f8cd43f080e0a0da53a97faf": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_acd02b4a31584748a6cdab140cfe39c1",
            "placeholder": "​",
            "style": "IPY_MODEL_3b525af8172b4da2bdd6b927a0edaba4",
            "value": "Loading pipeline components...: 100%"
          }
        },
        "b7c5f4ed54bb4cfa98d402302891ac4f": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f83abbc20e95491eb5e730651ff8d45e",
            "max": 2,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_a15a91d12e094d7c91d1db2337bbeb75",
            "value": 2
          }
        },
        "fa98e160e55146ba86ceb9aa137e07f7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_4ea0c598aaa04fabbede270e2a165efb",
            "placeholder": "​",
            "style": "IPY_MODEL_1e72c6cc511f43849bab0c4128a89c55",
            "value": " 2/2 [00:00&lt;00:00,  5.20it/s]"
          }
        },
        "2bcea0ab312b4ea8a474bebea04edeae": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "acd02b4a31584748a6cdab140cfe39c1": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3b525af8172b4da2bdd6b927a0edaba4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "f83abbc20e95491eb5e730651ff8d45e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "a15a91d12e094d7c91d1db2337bbeb75": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "4ea0c598aaa04fabbede270e2a165efb": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "1e72c6cc511f43849bab0c4128a89c55": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/fjadidi2001/Image_Inpaint/blob/main/Image_Inpaint_CelebHQ_Diffusions.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install diffusers transformers datasets accelerate safetensors"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ckYiBp8og8i7",
        "outputId": "510f2759-6762-4805-83c0-7bd1fd6afb97"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: diffusers in /usr/local/lib/python3.11/dist-packages (0.32.2)\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.11/dist-packages (4.47.1)\n",
            "Collecting datasets\n",
            "  Downloading datasets-3.2.0-py3-none-any.whl.metadata (20 kB)\n",
            "Requirement already satisfied: accelerate in /usr/local/lib/python3.11/dist-packages (1.2.1)\n",
            "Requirement already satisfied: safetensors in /usr/local/lib/python3.11/dist-packages (0.5.2)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.11/dist-packages (from diffusers) (8.6.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from diffusers) (3.17.0)\n",
            "Requirement already satisfied: huggingface-hub>=0.23.2 in /usr/local/lib/python3.11/dist-packages (from diffusers) (0.27.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.11/dist-packages (from diffusers) (1.26.4)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from diffusers) (2024.11.6)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.11/dist-packages (from diffusers) (2.32.3)\n",
            "Requirement already satisfied: Pillow in /usr/local/lib/python3.11/dist-packages (from diffusers) (11.1.0)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.11/dist-packages (from transformers) (24.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.11/dist-packages (from transformers) (6.0.2)\n",
            "Requirement already satisfied: tokenizers<0.22,>=0.21 in /usr/local/lib/python3.11/dist-packages (from transformers) (0.21.0)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.11/dist-packages (from transformers) (4.67.1)\n",
            "Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.11/dist-packages (from datasets) (17.0.0)\n",
            "Collecting dill<0.3.9,>=0.3.0 (from datasets)\n",
            "  Downloading dill-0.3.8-py3-none-any.whl.metadata (10 kB)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.11/dist-packages (from datasets) (2.2.2)\n",
            "Collecting xxhash (from datasets)\n",
            "  Downloading xxhash-3.5.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (12 kB)\n",
            "Collecting multiprocess<0.70.17 (from datasets)\n",
            "  Downloading multiprocess-0.70.16-py311-none-any.whl.metadata (7.2 kB)\n",
            "Collecting fsspec<=2024.9.0,>=2023.1.0 (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets)\n",
            "  Downloading fsspec-2024.9.0-py3-none-any.whl.metadata (11 kB)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.11/dist-packages (from datasets) (3.11.11)\n",
            "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from accelerate) (5.9.5)\n",
            "Requirement already satisfied: torch>=1.10.0 in /usr/local/lib/python3.11/dist-packages (from accelerate) (2.5.1+cu121)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (2.4.4)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (1.3.2)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (24.3.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (1.5.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (6.1.0)\n",
            "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (0.2.1)\n",
            "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets) (1.18.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub>=0.23.2->diffusers) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests->diffusers) (3.4.1)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.11/dist-packages (from requests->diffusers) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests->diffusers) (2.3.0)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.11/dist-packages (from requests->diffusers) (2024.12.14)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (3.4.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (3.1.5)\n",
            "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (12.1.105)\n",
            "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (9.1.0.70)\n",
            "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (12.1.3.1)\n",
            "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (11.0.2.54)\n",
            "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (10.3.2.106)\n",
            "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (11.4.5.107)\n",
            "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (12.1.0.106)\n",
            "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (2.21.5)\n",
            "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (12.1.105)\n",
            "Requirement already satisfied: triton==3.1.0 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (3.1.0)\n",
            "Requirement already satisfied: sympy==1.13.1 in /usr/local/lib/python3.11/dist-packages (from torch>=1.10.0->accelerate) (1.13.1)\n",
            "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.11/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.10.0->accelerate) (12.6.85)\n",
            "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /usr/local/lib/python3.11/dist-packages (from sympy==1.13.1->torch>=1.10.0->accelerate) (1.3.0)\n",
            "Requirement already satisfied: zipp>=3.20 in /usr/local/lib/python3.11/dist-packages (from importlib-metadata->diffusers) (3.21.0)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets) (2025.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.11/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.17.0)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch>=1.10.0->accelerate) (3.0.2)\n",
            "Downloading datasets-3.2.0-py3-none-any.whl (480 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m480.6/480.6 kB\u001b[0m \u001b[31m15.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading dill-0.3.8-py3-none-any.whl (116 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m8.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading fsspec-2024.9.0-py3-none-any.whl (179 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m179.3/179.3 kB\u001b[0m \u001b[31m11.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading multiprocess-0.70.16-py311-none-any.whl (143 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m143.5/143.5 kB\u001b[0m \u001b[31m12.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading xxhash-3.5.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (194 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.8/194.8 kB\u001b[0m \u001b[31m16.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: xxhash, fsspec, dill, multiprocess, datasets\n",
            "  Attempting uninstall: fsspec\n",
            "    Found existing installation: fsspec 2024.10.0\n",
            "    Uninstalling fsspec-2024.10.0:\n",
            "      Successfully uninstalled fsspec-2024.10.0\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "gcsfs 2024.10.0 requires fsspec==2024.10.0, but you have fsspec 2024.9.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed datasets-3.2.0 dill-0.3.8 fsspec-2024.9.0 multiprocess-0.70.16 xxhash-3.5.0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "from torchvision import transforms\n",
        "import torch\n",
        "import numpy as np\n",
        "import cv2\n",
        "from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel\n",
        "from torch import nn\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.optim import Adam\n",
        "import gc\n",
        "\n",
        "# --- Memory Optimization Setup ---\n",
        "torch.backends.cudnn.benchmark = True\n",
        "torch.cuda.empty_cache()\n",
        "\n",
        "# --- Reduced Batch Size ---\n",
        "BATCH_SIZE = 2\n",
        "\n",
        "# --- Dataset with Smaller Resolution ---\n",
        "transform = transforms.Compose([\n",
        "    transforms.Resize((128, 128)),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize([0.5], [0.5])\n",
        "])\n",
        "\n",
        "# --- Corrected Mask Generation ---\n",
        "def create_irregular_mask(image_size=(128, 128)):\n",
        "    H, W = image_size\n",
        "    mask = np.zeros((H, W), dtype=np.uint8)\n",
        "    x = np.random.randint(0, W//2)\n",
        "    y = np.random.randint(0, H//2)\n",
        "    w = np.random.randint(W//4, W - x)\n",
        "    h = np.random.randint(H//4, H - y)\n",
        "    cv2.rectangle(mask, (x, y), (x + w, y + h), 1, -1)\n",
        "    return mask\n",
        "\n",
        "# --- Model Loading with Explicit safetensors Disable ---\n",
        "try:\n",
        "    # First try loading with safetensors disabled\n",
        "    model = DDPMPipeline.from_pretrained(\n",
        "        \"google/ddpm-celebahq-256\",\n",
        "        use_safetensors=False,\n",
        "        safety_checker=None\n",
        "    )\n",
        "except Exception as e:\n",
        "    print(f\"Error loading model: {e}\")\n",
        "    # Fallback to basic UNet architecture\n",
        "    model = DDPMPipeline.from_pretrained(\n",
        "        UNet2DModel(\n",
        "            sample_size=128,\n",
        "            in_channels=3,\n",
        "            out_channels=3,\n",
        "            layers_per_block=2,\n",
        "            block_out_channels=(128, 256, 512),\n",
        "            down_block_types=(\n",
        "                \"DownBlock2D\",\n",
        "                \"DownBlock2D\",\n",
        "                \"DownBlock2D\"\n",
        "            ),\n",
        "            up_block_types=(\n",
        "                \"UpBlock2D\",\n",
        "                \"UpBlock2D\",\n",
        "                \"UpBlock2D\"\n",
        "            ),\n",
        "        ),\n",
        "        safety_checker=None\n",
        "    )\n",
        "\n",
        "# --- Modified UNet with Proper Initialization ---\n",
        "original_conv_in = model.unet.conv_in\n",
        "\n",
        "# Create new conv layer with 4 input channels\n",
        "new_conv_in = nn.Conv2d(4, original_conv_in.out_channels,\n",
        "                       kernel_size=original_conv_in.kernel_size,\n",
        "                       padding=original_conv_in.padding,\n",
        "                       stride=original_conv_in.stride)\n",
        "\n",
        "with torch.no_grad():\n",
        "    # Handle weights\n",
        "    if original_conv_in.weight is not None:\n",
        "        new_conv_in.weight[:, :3] = original_conv_in.weight.clone()\n",
        "        new_conv_in.weight[:, 3:] = 0  # Initialize mask channel to zero\n",
        "    else:\n",
        "        new_conv_in.weight = nn.Parameter(torch.randn_like(new_conv_in.weight))\n",
        "\n",
        "    # Handle bias with proper Parameter conversion\n",
        "    if original_conv_in.bias is not None:\n",
        "        new_conv_in.bias = nn.Parameter(original_conv_in.bias.clone())\n",
        "    else:\n",
        "        new_conv_in.bias = None  # Explicitly set to None if no bias\n",
        "\n",
        "model.unet.conv_in = new_conv_in\n",
        "\n",
        "# --- Device Setup ---\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model.unet.to(device)\n",
        "model.unet.train()\n",
        "\n",
        "# --- Data Loading ---\n",
        "ds = load_dataset(\"saitsharipov/CelebA-HQ\")\n",
        "\n",
        "def collate_fn(batch):\n",
        "    images = [transform(img['image'].convert('RGB')) for img in batch]\n",
        "    return torch.stack(images)\n",
        "\n",
        "\n",
        "train_loader = DataLoader(\n",
        "    ds['train'],\n",
        "    batch_size=BATCH_SIZE,\n",
        "    shuffle=True,\n",
        "    collate_fn=collate_fn,\n",
        "    num_workers=0,  # Set num_workers to 0 to disable multiprocessing\n",
        "    pin_memory=True\n",
        ")\n",
        "\n",
        "# --- Training Setup ---\n",
        "optimizer = Adam(model.unet.parameters(), lr=1e-4)\n",
        "\n",
        "\n",
        "# --- Training Loop ---\n",
        "for epoch in range(5):\n",
        "    for images in train_loader:\n",
        "        # Move images and masks to the device inside the training loop\n",
        "        images = images.to(device, non_blocking=True)\n",
        "        masks = torch.stack([\n",
        "            torch.from_numpy(create_irregular_mask()).float()\n",
        "            for _ in range(images.size(0))\n",
        "        ]).unsqueeze(1).to(device, non_blocking=True)\n",
        "\n",
        "        masked_images = images * (1 - masks)\n",
        "        noise = torch.randn_like(images)\n",
        "        timesteps = torch.randint(0, model.scheduler.config.num_train_timesteps,\n",
        "                                (BATCH_SIZE,), device=device)\n",
        "\n",
        "        with torch.cuda.amp.autocast():\n",
        "            noisy_images = model.scheduler.add_noise(images, noise, timesteps)\n",
        "            model_output = model.unet(\n",
        "                torch.cat([noisy_images, masks], dim=1),\n",
        "                timesteps\n",
        "            ).sample\n",
        "            loss = nn.functional.mse_loss(model_output, noise)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(model.unet.parameters(), 1.0)\n",
        "        optimizer.step()\n",
        "\n",
        "        del noisy_images, masks, model_output\n",
        "        torch.cuda.empty_cache()\n",
        "\n",
        "    print(f\"Epoch {epoch+1} Loss: {loss.item():.4f}\")\n",
        "\n",
        "# Save model\n",
        "model.save_pretrained(\"inpainting_model\", safe_serialization=torch.cuda.is_available())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 121,
          "referenced_widgets": [
            "e9a7aeae5a474eab8f60ee3046d7cebc",
            "c3ec9930f8cd43f080e0a0da53a97faf",
            "b7c5f4ed54bb4cfa98d402302891ac4f",
            "fa98e160e55146ba86ceb9aa137e07f7",
            "2bcea0ab312b4ea8a474bebea04edeae",
            "acd02b4a31584748a6cdab140cfe39c1",
            "3b525af8172b4da2bdd6b927a0edaba4",
            "f83abbc20e95491eb5e730651ff8d45e",
            "a15a91d12e094d7c91d1db2337bbeb75",
            "4ea0c598aaa04fabbede270e2a165efb",
            "1e72c6cc511f43849bab0c4128a89c55"
          ]
        },
        "id": "7o3R4yb6hD6k",
        "outputId": "eac4aded-ca04-4518-db27-60482b0773b5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Keyword arguments {'safety_checker': None} are not expected by DDPMPipeline and will be ignored.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Loading pipeline components...:   0%|          | 0/2 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "e9a7aeae5a474eab8f60ee3046d7cebc"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "<ipython-input-11-e0a373504886>:135: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
            "  with torch.cuda.amp.autocast():\n"
          ]
        }
      ]
    }
  ]
}