{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"Before you judge me for creating a stable diffusion notebook when GANs were specifically required in this competition, I have just two things to say. First, when this competition was launched, stable diffusion wasn't a thing. Secondly, this notebook is for learning and experimenting with SD models, which are clearly the industry standard for generating images.\n\n**Disclaimer**: *No GANs were harmed during the training of this model* 😛","metadata":{}},{"cell_type":"markdown","source":"Stable diffusion operating on image space are compute expensive, and take longer to train. If you're planning to train on high res images, you might run out of memory. This problem is overcom by latent diffusion models.\n\nInstead of passing the entire image into the unet, we pass a latent tensor to it. This tensor, generally much smaller in size than the original image, is generated as an intermediate activations to a variational autoencoder.","metadata":{}},{"cell_type":"code","source":"#pass vae.encode instead of vae to accelerate","metadata":{"execution":{"iopub.status.busy":"2023-09-25T09:22:52.461344Z","iopub.execute_input":"2023-09-25T09:22:52.462035Z","iopub.status.idle":"2023-09-25T09:22:52.469498Z","shell.execute_reply.started":"2023-09-25T09:22:52.461991Z","shell.execute_reply":"2023-09-25T09:22:52.468276Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Imports","metadata":{}},{"cell_type":"markdown","source":"I'll be using huggingface diffusers library to train my model","metadata":{}},{"cell_type":"code","source":"!pip install git+https://github.com/huggingface/accelerate -q\n!pip install git+https://github.com/huggingface/diffusers -q","metadata":{"execution":{"iopub.status.busy":"2023-09-25T09:22:52.474249Z","iopub.execute_input":"2023-09-25T09:22:52.475219Z","iopub.status.idle":"2023-09-25T09:23:54.027019Z","shell.execute_reply.started":"2023-09-25T09:22:52.475183Z","shell.execute_reply":"2023-09-25T09:23:54.025674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from diffusers import AutoencoderKL,DDPMScheduler, UNet2DModel\nfrom PIL import Image\nimport torchvision.transforms as transforms\nimport torch\nimport numpy as np\nimport matplotlib.pyplot as plt","metadata":{"execution":{"iopub.status.busy":"2023-09-25T09:23:54.029429Z","iopub.execute_input":"2023-09-25T09:23:54.029917Z","iopub.status.idle":"2023-09-25T09:24:14.835974Z","shell.execute_reply.started":"2023-09-25T09:23:54.029878Z","shell.execute_reply":"2023-09-25T09:24:14.835008Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Connection with Huggingface hub\n\nIn order to share your model on the hub, you’ll need to login to your Hugging Face account. Run the cell below and pass your access token to login. **This model will be public and anybody can use it in their notebooks**.\n\nLink to my Huggingface repo: https://huggingface.co/NavneetSajwan/unconditional-monet-128","metadata":{}},{"cell_type":"code","source":"from huggingface_hub import notebook_login\nnotebook_login()","metadata":{"execution":{"iopub.status.busy":"2023-09-25T09:24:14.837373Z","iopub.execute_input":"2023-09-25T09:24:14.838113Z","iopub.status.idle":"2023-09-25T09:24:14.869478Z","shell.execute_reply.started":"2023-09-25T09:24:14.838075Z","shell.execute_reply":"2023-09-25T09:24:14.868199Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Since the model checkpoints are quite large, install Git-LFS to version these large files:","metadata":{}},{"cell_type":"code","source":"!sudo apt -qq install git-lfs\n!git config --global credential.helper store","metadata":{"execution":{"iopub.status.busy":"2023-09-25T09:24:14.871632Z","iopub.execute_input":"2023-09-25T09:24:14.872052Z","iopub.status.idle":"2023-09-25T09:24:19.064797Z","shell.execute_reply.started":"2023-09-25T09:24:14.871985Z","shell.execute_reply":"2023-09-25T09:24:19.063459Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Preprocessing","metadata":{}},{"cell_type":"markdown","source":"### Utility functions","metadata":{}},{"cell_type":"code","source":"def tensors_to_latents(input_im, vae):\n    #batch to latents\n    with torch.no_grad():\n        latent = vae.encode((input_im)*2-1) # Note scaling\n    return 0.18215 * latent.latent_dist.sample()  \n\ndef pil_to_latent(input_im, vae):\n    # Single image -> single latent in a batch (so size 1, 4, 64, 64)\n    with torch.no_grad():\n        latent = vae.encode(transforms.ToTensor()(input_im).unsqueeze(0)*2-1) # Note scaling\n    return 0.18215 * latent.latent_dist.sample()\n\ndef latents_to_pil(latents,vae):\n    # bath of latents -> list of images\n    latents = (1 / 0.18215) * latents\n    with torch.no_grad():\n        image = vae.decode(latents).sample\n    image = (image / 2 + 0.5).clamp(0, 1)\n    image = image.detach().cpu().permute(0, 2, 3, 1).numpy()\n    images = (image * 255).round().astype(\"uint8\")\n    pil_images = [Image.fromarray(image) for image in images]\n    return pil_images","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:24:15.560950Z","iopub.execute_input":"2023-09-24T15:24:15.561279Z","iopub.status.idle":"2023-09-24T15:24:15.570378Z","shell.execute_reply.started":"2023-09-24T15:24:15.561248Z","shell.execute_reply":"2023-09-24T15:24:15.569515Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def pil_to_tensor(pil_image):\n    # Define a transformation to convert the PIL Image to a tensor\n    transform = transforms.ToTensor()\n\n    # Apply the transformation to the image\n    image_tensor = transform(pil_image)\n\n    # Check the shape of the resulting tensor\n    print(image_tensor.shape)\n    return image_tensor\n\ndef tensor_to_pil_image(tensor):\n    # Make sure the tensor has the expected shape [1, 3, 256, 256]\n    if tensor.shape != (1, 3, 256, 256):\n        raise ValueError(\"Input tensor must have shape [1, 3, 256, 256]\")\n\n    # Convert the tensor to a PIL Image\n    pil_image = transforms.ToPILImage()(tensor.squeeze(0).cpu())\n\n    return pil_image\n\ndef plot_feature_maps(latent_tensor):\n    # Make sure the tensor has the expected shape [1, 4, 32, 32]\n    if latent_tensor.shape != (1, 4, 32, 32):\n        raise ValueError(\"Input tensor must have shape [1, 4, 32, 32]\")\n\n    # Convert the tensor to a NumPy array and squeeze the batch dimension\n    feature_maps = latent_tensor.squeeze(0).detach().cpu().numpy()\n\n    # Get the number of feature maps\n    num_feature_maps = feature_maps.shape[0]\n\n    # Create a grid of subplots\n    rows = 1\n    cols = num_feature_maps\n    fig, axes = plt.subplots(rows, cols, figsize=(12, 4))\n\n    # Plot each feature map\n    for i in range(num_feature_maps):\n        ax = axes[i] if num_feature_maps > 1 else axes  # Use single subplot if there's only one feature map\n        ax.imshow(feature_maps[i], cmap='gray')  # You can change the colormap as needed\n        ax.set_title(f\"Feature Map {i+1}\")\n        ax.axis('off')\n\n    plt.tight_layout()\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:24:15.571990Z","iopub.execute_input":"2023-09-24T15:24:15.572628Z","iopub.status.idle":"2023-09-24T15:24:15.586008Z","shell.execute_reply.started":"2023-09-24T15:24:15.572596Z","shell.execute_reply":"2023-09-24T15:24:15.585130Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### VAE","metadata":{}},{"cell_type":"code","source":"#get autoencoder\nvae = AutoencoderKL.from_pretrained(\"CompVis/stable-diffusion-v1-1\",\n                                    subfolder = \"vae\",\n                                    safe_tensors=True\n                                   )","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:24:15.587885Z","iopub.execute_input":"2023-09-24T15:24:15.588290Z","iopub.status.idle":"2023-09-24T15:24:23.099239Z","shell.execute_reply.started":"2023-09-24T15:24:15.588258Z","shell.execute_reply":"2023-09-24T15:24:23.098206Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimgsdir = \"/kaggle/input/gan-getting-started/monet_jpg/\"\nfnamelist = os.listdir(f\"{imgsdir}\")","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:24:23.102551Z","iopub.execute_input":"2023-09-24T15:24:23.102906Z","iopub.status.idle":"2023-09-24T15:24:23.136829Z","shell.execute_reply.started":"2023-09-24T15:24:23.102872Z","shell.execute_reply":"2023-09-24T15:24:23.135978Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"i = 10\nfpath = os.path.join(imgsdir, fnamelist[i])\nfpath","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:24:23.138155Z","iopub.execute_input":"2023-09-24T15:24:23.138697Z","iopub.status.idle":"2023-09-24T15:24:23.146718Z","shell.execute_reply.started":"2023-09-24T15:24:23.138664Z","shell.execute_reply":"2023-09-24T15:24:23.145833Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample = Image.open(fpath)\nsample","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:24:23.147973Z","iopub.execute_input":"2023-09-24T15:24:23.148952Z","iopub.status.idle":"2023-09-24T15:24:23.223602Z","shell.execute_reply.started":"2023-09-24T15:24:23.148908Z","shell.execute_reply":"2023-09-24T15:24:23.222818Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"latents = pil_to_latent(sample,vae)\nlatents.shape","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:06.645193Z","iopub.execute_input":"2023-09-24T15:25:06.645885Z","iopub.status.idle":"2023-09-24T15:25:09.929657Z","shell.execute_reply.started":"2023-09-24T15:25:06.645803Z","shell.execute_reply":"2023-09-24T15:25:09.928522Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport torch\n\n# Assuming latent_tensor is a tensor of shape [1, 4, 32, 32]\nplot_feature_maps(latents)\n","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:09.931800Z","iopub.execute_input":"2023-09-24T15:25:09.932173Z","iopub.status.idle":"2023-09-24T15:25:10.589828Z","shell.execute_reply.started":"2023-09-24T15:25:09.932139Z","shell.execute_reply":"2023-09-24T15:25:10.588722Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pilimg = latents_to_pil(latents,vae)\npilimg[0]","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:10.591662Z","iopub.execute_input":"2023-09-24T15:25:10.592343Z","iopub.status.idle":"2023-09-24T15:25:18.558769Z","shell.execute_reply.started":"2023-09-24T15:25:10.592303Z","shell.execute_reply":"2023-09-24T15:25:18.557686Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"noise_scheduler = DDPMScheduler(num_train_timesteps=1000)\nnoise = torch.randn(latents.shape)\ntimesteps = torch.LongTensor([150])\nnoisy_latent = noise_scheduler.add_noise(latents, noise, timesteps)\n# uses 100 to add noise from the 1000 steps\n# indexes at 100, to get a noise value","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:18.561417Z","iopub.execute_input":"2023-09-24T15:25:18.561776Z","iopub.status.idle":"2023-09-24T15:25:18.587662Z","shell.execute_reply.started":"2023-09-24T15:25:18.561731Z","shell.execute_reply":"2023-09-24T15:25:18.586676Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pilimg = latents_to_pil(noisy_latent,vae)\npilimg[0]","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:18.588976Z","iopub.execute_input":"2023-09-24T15:25:18.589837Z","iopub.status.idle":"2023-09-24T15:25:25.541870Z","shell.execute_reply.started":"2023-09-24T15:25:18.589799Z","shell.execute_reply":"2023-09-24T15:25:25.540682Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#For convenience, create a TrainingConfig class containing the training hyperparameters \n\nfrom dataclasses import dataclass\n\n@dataclass\nclass TrainingConfig:\n    image_size = 32  # the generated image resolution\n    train_batch_size = 2\n    eval_batch_size = 2  # how many images to sample during evaluation\n    num_epochs = 10\n    gradient_accumulation_steps = 1\n    learning_rate = 1e-4\n    lr_warmup_steps = 500\n    save_image_epochs = 10\n    save_model_epochs = 50\n    mixed_precision = \"fp16\"  # `no` for float32, `fp16` for automatic mixed precision\n    output_dir = \"latents-trained-monets-256\"  # the model name locally and on the HF Hub\n\n    push_to_hub = True  # whether to upload the saved model to the HF Hub\n    hub_private_repo = False\n    overwrite_output_dir = True  # overwrite the old model when re-running the notebook\n    seed = 0\n\nconfig = TrainingConfig()","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:25.543340Z","iopub.execute_input":"2023-09-24T15:25:25.543996Z","iopub.status.idle":"2023-09-24T15:25:25.551989Z","shell.execute_reply.started":"2023-09-24T15:25:25.543961Z","shell.execute_reply":"2023-09-24T15:25:25.550926Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = UNet2DModel(\n    sample_size=config.image_size,  # the target image resolution\n    in_channels=4,  # the number of input channels, 3 for RGB images\n    out_channels=4,  # the number of output channels\n    layers_per_block=2,  # how many ResNet layers to use per UNet block\n    block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n    down_block_types=(\n        \"DownBlock2D\",  # a regular ResNet downsampling block\n        \"DownBlock2D\",\n        \"DownBlock2D\",\n        \"DownBlock2D\",\n        \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n        \"DownBlock2D\",\n    ),\n    up_block_types=(\n        \"UpBlock2D\",  # a regular ResNet upsampling block\n        \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n        \"UpBlock2D\",\n        \"UpBlock2D\",\n        \"UpBlock2D\",\n        \"UpBlock2D\",\n    ),\n)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:25.553391Z","iopub.execute_input":"2023-09-24T15:25:25.553991Z","iopub.status.idle":"2023-09-24T15:25:26.511922Z","shell.execute_reply.started":"2023-09-24T15:25:25.553957Z","shell.execute_reply":"2023-09-24T15:25:26.510961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from datasets import load_dataset\n\n# Load the monet jpgs into huggingface Dataset object\ndataset = load_dataset(\"imagefolder\", data_dir = \"/kaggle/input/gan-getting-started/monet_jpg\", split=\"train\")","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:26.513539Z","iopub.execute_input":"2023-09-24T15:25:26.514195Z","iopub.status.idle":"2023-09-24T15:25:31.253466Z","shell.execute_reply.started":"2023-09-24T15:25:26.514158Z","shell.execute_reply":"2023-09-24T15:25:31.252391Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"It's evident that num_rows represents the number of images present in dataset","metadata":{}},{"cell_type":"markdown","source":"Says, we have PIL images of size 256 by 256. We'll use progressive learning technique and train the model on 128 by 128 images first. Then we'll use this trained model and improve it using 256 by 256 input images ","metadata":{}},{"cell_type":"code","source":"#display a minibatch\nimport matplotlib.pyplot as plt\n\nfig, axs = plt.subplots(1, 4, figsize=(16, 4))\nfor i, image in enumerate(dataset[:4][\"image\"]):\n    axs[i].imshow(image)\n    axs[i].set_axis_off()\nfig.show()","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:31.255051Z","iopub.execute_input":"2023-09-24T15:25:31.256095Z","iopub.status.idle":"2023-09-24T15:25:31.761524Z","shell.execute_reply.started":"2023-09-24T15:25:31.256056Z","shell.execute_reply":"2023-09-24T15:25:31.760693Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"The images are all different sizes though, so you’ll need to preprocess them first:\n\n- Resize changes the image size to the one defined in config.image_size.\n- RandomHorizontalFlip augments the dataset by randomly mirroring the images.\n- Normalize is important to rescale the pixel values into a [-1, 1] range, which is what the model expects.","metadata":{}},{"cell_type":"code","source":"from torchvision import transforms\n\npreprocess = transforms.Compose(\n    [\n        transforms.Resize((256, 256)),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToTensor(),\n        transforms.Normalize([0.5], [0.5]),\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:31.766555Z","iopub.execute_input":"2023-09-24T15:25:31.767381Z","iopub.status.idle":"2023-09-24T15:25:31.773493Z","shell.execute_reply.started":"2023-09-24T15:25:31.767343Z","shell.execute_reply":"2023-09-24T15:25:31.772495Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Use 🤗 Datasets’ set_transform method to apply the preprocess function on the fly during training:","metadata":{}},{"cell_type":"code","source":"def transform(examples):\n    images = [preprocess(image.convert(\"RGB\")) for image in examples[\"image\"]]\n    return {\"images\": images}\n\ndataset.set_transform(transform)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:31.774708Z","iopub.execute_input":"2023-09-24T15:25:31.775827Z","iopub.status.idle":"2023-09-24T15:25:31.791245Z","shell.execute_reply.started":"2023-09-24T15:25:31.775793Z","shell.execute_reply":"2023-09-24T15:25:31.790102Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config.train_batch_size","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:31.792730Z","iopub.execute_input":"2023-09-24T15:25:31.793094Z","iopub.status.idle":"2023-09-24T15:25:31.806804Z","shell.execute_reply.started":"2023-09-24T15:25:31.793062Z","shell.execute_reply":"2023-09-24T15:25:31.805653Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\ntrain_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:31.808543Z","iopub.execute_input":"2023-09-24T15:25:31.809149Z","iopub.status.idle":"2023-09-24T15:25:31.820916Z","shell.execute_reply.started":"2023-09-24T15:25:31.809115Z","shell.execute_reply":"2023-09-24T15:25:31.819990Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create a UNet2DModel\nPretrained models in 🧨 Diffusers are easily created from their model class with the parameters you want. For example, to create a UNet2DModel:","metadata":{}},{"cell_type":"code","source":"# check if the input and and output are of same size\ntorch_img = dataset[0][\"images\"].unsqueeze(0)\ntorch_img.shape","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:31.822314Z","iopub.execute_input":"2023-09-24T15:25:31.823007Z","iopub.status.idle":"2023-09-24T15:25:31.849213Z","shell.execute_reply.started":"2023-09-24T15:25:31.822972Z","shell.execute_reply":"2023-09-24T15:25:31.848203Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# torch tensor to latent\nlatent = tensors_to_latents(torch_img,vae)\nlatent.shape","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:31.850660Z","iopub.execute_input":"2023-09-24T15:25:31.851189Z","iopub.status.idle":"2023-09-24T15:25:35.026862Z","shell.execute_reply.started":"2023-09-24T15:25:31.851156Z","shell.execute_reply":"2023-09-24T15:25:35.025708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model = UNet2DModel(\n#     sample_size=config.image_size,  # the target image resolution\n#     in_channels=4,  # the number of input channels, 3 for RGB images\n#     out_channels=4,  # the number of output channels\n#     layers_per_block=3,  # how many ResNet layers to use per UNet block\n# #     block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n#     block_out_channels=(128, 128, 256, 256, 512,512),  # the number of output channels for each UNet block\n#     down_block_types=(\n#         \"DownBlock2D\",  # a regular ResNet downsampling block\n#         \"DownBlock2D\",\n#         \"DownBlock2D\",\n#         \"DownBlock2D\",\n#         \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n#         \"DownBlock2D\",\n#     ),\n#     up_block_types=(\n#         \"UpBlock2D\",  # a regular ResNet upsampling block\n#         \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n#         \"UpBlock2D\",\n#         \"UpBlock2D\",\n#         \"UpBlock2D\",\n#         \"UpBlock2D\",\n#     ),\n# )","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:35.029868Z","iopub.execute_input":"2023-09-24T15:25:35.030159Z","iopub.status.idle":"2023-09-24T15:25:35.037530Z","shell.execute_reply.started":"2023-09-24T15:25:35.030133Z","shell.execute_reply":"2023-09-24T15:25:35.036641Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# model = UNet2DModel(\n#     sample_size=config.image_size,  # the target image resolution\n#     in_channels=4,  # the number of input channels, 3 for RGB images\n#     out_channels=4,  # the number of output channels\n#     layers_per_block=3,  # how many ResNet layers to use per UNet block\n# #     block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n#     block_out_channels=(128, 256, 256, 512,512),  # the number of output channels for each UNet block\n#     down_block_types=(\n#         \"DownBlock2D\",  # a regular ResNet downsampling block\n#         \"DownBlock2D\",\n#         \"DownBlock2D\",\n# #         \"DownBlock2D\",\n#         \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n#         \"DownBlock2D\",\n#     ),\n#     up_block_types=(\n#         \"UpBlock2D\",  # a regular ResNet upsampling block\n#         \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n#         \"UpBlock2D\",\n#         \"UpBlock2D\",\n# #         \"UpBlock2D\",\n#         \"UpBlock2D\",\n#     ),\n# )\n\n# def fw_hook(module, input, output):\n#     if torch.is_tensor(output):\n#         print(f\"{output.shape}\")\n#     else:\n#         print(\"output not a tensor\")\n#     print(\"==========================\")\n\n# for name, layer in model.named_modules():\n#     layer.register_forward_hook(fw_hook)\n","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:35.040588Z","iopub.execute_input":"2023-09-24T15:25:35.040915Z","iopub.status.idle":"2023-09-24T15:25:35.051598Z","shell.execute_reply.started":"2023-09-24T15:25:35.040888Z","shell.execute_reply":"2023-09-24T15:25:35.050737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"Input shape:\", latents.shape)\n\nprint(\"Output shape:\", model(latents, timestep=0).sample.shape)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:01:29.510037Z","iopub.execute_input":"2023-09-24T16:01:29.510429Z","iopub.status.idle":"2023-09-24T16:01:30.057538Z","shell.execute_reply.started":"2023-09-24T16:01:29.510395Z","shell.execute_reply":"2023-09-24T16:01:30.056410Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Create a scheduler\nThe scheduler behaves differently depending on whether you’re using the model for training or inference. During inference, the scheduler generates image from the noise. During training, the scheduler takes a model output - or a sample - from a specific point in the diffusion process and applies noise to the image according to a noise schedule and an update rule.\n\nLet’s take a look at the DDPMScheduler and use the add_noise method to add some random noise to the sample_image from before:\n\nCopied\n","metadata":{}},{"cell_type":"code","source":"\nnoise_scheduler = DDPMScheduler(num_train_timesteps=1000)\nnoise = torch.randn(latents.shape)\ntimesteps = torch.LongTensor([50])\nnoisy_latent = noise_scheduler.add_noise(latents, noise, timesteps)\n\nImage.fromarray(((noisy_latent.permute(0, 2, 3, 1) + 1.0) * 127.5).type(torch.uint8).cpu().numpy()[0])","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:35.314704Z","iopub.execute_input":"2023-09-24T15:25:35.315036Z","iopub.status.idle":"2023-09-24T15:25:35.326919Z","shell.execute_reply.started":"2023-09-24T15:25:35.315006Z","shell.execute_reply":"2023-09-24T15:25:35.325912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"noisy_latent.shape, timesteps.shape","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:35.328367Z","iopub.execute_input":"2023-09-24T15:25:35.328991Z","iopub.status.idle":"2023-09-24T15:25:35.337646Z","shell.execute_reply.started":"2023-09-24T15:25:35.328958Z","shell.execute_reply":"2023-09-24T15:25:35.336775Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.nn.functional as F\n# The training objective of the model is to predict the noise added to the image. \n# The loss at this step can be calculated by:\nnoise_pred = model(noisy_latent, timesteps).sample\n# noise_pred.shape, noise.shape\nloss = F.mse_loss(noise_pred, noise)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:35.338741Z","iopub.execute_input":"2023-09-24T15:25:35.339711Z","iopub.status.idle":"2023-09-24T15:25:35.549939Z","shell.execute_reply.started":"2023-09-24T15:25:35.339685Z","shell.execute_reply":"2023-09-24T15:25:35.548929Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train the model","metadata":{}},{"cell_type":"code","source":"def fw_hook(module, input, output):\n    if torch.is_tensor(output):\n        print(f\"{output.shape}\")\n    else:\n        print(\"output not a tensor\")\n    print(\"==========================\")\n\nfor name, layer in model.named_modules():\n    layer.register_forward_hook(fw_hook)\n","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:01:23.935704Z","iopub.execute_input":"2023-09-24T16:01:23.936501Z","iopub.status.idle":"2023-09-24T16:01:23.943748Z","shell.execute_reply.started":"2023-09-24T16:01:23.936463Z","shell.execute_reply":"2023-09-24T16:01:23.942831Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = UNet2DModel(\n    sample_size=config.image_size,  # the target image resolution\n    in_channels=4,  # the number of input channels, 3 for RGB images\n    out_channels=4,  # the number of output channels\n    layers_per_block=4,  # how many ResNet layers to use per UNet block\n#     block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block\n    block_out_channels=(128, 256, 512,512),  # the number of output channels for each UNet block\n    down_block_types=(\n        \"DownBlock2D\",  # a regular ResNet downsampling block\n#         \"DownBlock2D\",\n        \"DownBlock2D\",\n#         \"DownBlock2D\",\n        \"AttnDownBlock2D\",  # a ResNet downsampling block with spatial self-attention\n        \"DownBlock2D\",\n    ),\n    up_block_types=(\n        \"UpBlock2D\",  # a regular ResNet upsampling block\n        \"AttnUpBlock2D\",  # a ResNet upsampling block with spatial self-attention\n#         \"UpBlock2D\",\n        \"UpBlock2D\",\n#         \"UpBlock2D\",\n        \"UpBlock2D\",\n    ),\n)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:02:00.443147Z","iopub.execute_input":"2023-09-24T16:02:00.443843Z","iopub.status.idle":"2023-09-24T16:02:01.816779Z","shell.execute_reply.started":"2023-09-24T16:02:00.443786Z","shell.execute_reply":"2023-09-24T16:02:01.815792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from diffusers.optimization import get_cosine_schedule_with_warmup\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)\nlr_scheduler = get_cosine_schedule_with_warmup(\n    optimizer=optimizer,\n    num_warmup_steps=config.lr_warmup_steps,\n    num_training_steps=(len(train_dataloader) * config.num_epochs),\n)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:02:05.278225Z","iopub.execute_input":"2023-09-24T16:02:05.278631Z","iopub.status.idle":"2023-09-24T16:02:05.288494Z","shell.execute_reply.started":"2023-09-24T16:02:05.278597Z","shell.execute_reply":"2023-09-24T16:02:05.286827Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# #get autoencoder\n# vae = AutoencoderKL.from_pretrained(\"CompVis/stable-diffusion-v1-1\",\n#                                     subfolder = \"vae\",\n#                                     safe_tensors=True\n#                                    )\n# evaluate(config, 1, ddpm, vae)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:25:35.563704Z","iopub.execute_input":"2023-09-24T15:25:35.564121Z","iopub.status.idle":"2023-09-24T15:25:35.575969Z","shell.execute_reply.started":"2023-09-24T15:25:35.564082Z","shell.execute_reply":"2023-09-24T15:25:35.575086Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from diffusers import DDPMPipeline\nfrom diffusers.utils import make_image_grid\nimport math\nimport os\n\n\ndef evaluate(config, epoch, pipeline, vae):\n    # Sample some images from random noise (this is the backward diffusion process).\n    # The default pipeline output type is `List[PIL.Image]`\n    images = pipeline(\n        num_inference_steps=40,\n        batch_size=config.eval_batch_size,\n        generator=torch.manual_seed(config.seed),\n        output_type=\"array\"\n    ).images # output as np array\n    print(images.shape)\n    # decode latents :(bs,h,w,c) [bs, c, h, w]\n    latents = torch.from_numpy(images.transpose(0, 3, 1, 2))\n    print(latents.shape)\n    print(latents.device, vae.device)\n    decodedimgs = latents_to_pil(latents.to(vae.device), vae)\n    print(\"decodedimgs:\", decodedimgs)\n#     decodedimgs = pil_to_tensor(decodedimgs)\n    # Make a grid out of the images\n    image_grid = make_image_grid(decodedimgs, rows=1, cols=2)\n\n    # Save the images\n    test_dir = os.path.join(config.output_dir, \"samples\")\n    os.makedirs(test_dir, exist_ok=True)\n    image_grid.save(f\"{test_dir}/{epoch:04d}.png\")","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:02:09.543739Z","iopub.execute_input":"2023-09-24T16:02:09.544801Z","iopub.status.idle":"2023-09-24T16:02:09.553904Z","shell.execute_reply.started":"2023-09-24T16:02:09.544742Z","shell.execute_reply":"2023-09-24T16:02:09.552724Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from accelerate import Accelerator\nfrom huggingface_hub import HfFolder, Repository, whoami\nfrom tqdm.auto import tqdm\nfrom pathlib import Path\nimport os\n\n\ndef get_full_repo_name(model_id: str, organization: str = None, token: str = None):\n    if token is None:\n        token = HfFolder.get_token()\n    if organization is None:\n        username = whoami(token)[\"name\"]\n        return f\"{username}/{model_id}\"\n    else:\n        return f\"{organization}/{model_id}\"\n\n\ndef train_loop(config, model, vae, noise_scheduler, optimizer, train_dataloader, lr_scheduler):\n    # Initialize accelerator and tensorboard logging\n    accelerator = Accelerator(\n        mixed_precision=config.mixed_precision,\n        gradient_accumulation_steps=config.gradient_accumulation_steps,\n        log_with=\"tensorboard\",\n        project_dir=os.path.join(config.output_dir, \"logs\"),\n    )\n    if accelerator.is_main_process:\n        if config.push_to_hub:\n            repo_name = get_full_repo_name(Path(config.output_dir).name)\n            repo = Repository(config.output_dir, clone_from=repo_name)\n        elif config.output_dir is not None:\n            os.makedirs(config.output_dir, exist_ok=True)\n        accelerator.init_trackers(\"train_example\")\n\n    # Prepare everything\n    # There is no specific order to remember, you just need to unpack the\n    # objects in the same order you gave them to the prepare method.\n    model, vae, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n        model, vae, optimizer, train_dataloader, lr_scheduler\n    )\n    print(model.device)\n    global_step = 0\n\n    # Now you train the model\n    for epoch in range(config.num_epochs):\n        progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)\n        progress_bar.set_description(f\"Epoch {epoch}\")\n\n        for step, batch in enumerate(train_dataloader):\n            clean_images = batch[\"images\"]\n            \n            clean_images = vae.encode(clean_images).latent_dist.sample()\n#             print(\"clean_images.device\")\n            \n            # Sample noise to add to the images\n            noise = torch.randn(clean_images.shape).to(clean_images.device)\n            bs = clean_images.shape[0]\n\n            # Sample a random timestep for each image\n            timesteps = torch.randint(\n                0, \n                noise_scheduler.config.num_train_timesteps//2, # using less noisy images\n                (bs,), \n                device=clean_images.device\n            ).long()\n\n            # Add noise to the clean images according to the noise magnitude at each timestep\n            # (this is the forward diffusion process)\n            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)\n\n            with accelerator.accumulate(model):\n                # Predict the noise residual\n                noise_pred = model(noisy_images, timesteps, return_dict=False)[0]\n                loss = F.mse_loss(noise_pred, noise)\n                accelerator.backward(loss)\n\n                accelerator.clip_grad_norm_(model.parameters(), 1.0)\n                optimizer.step()\n                lr_scheduler.step()\n                optimizer.zero_grad()\n\n            progress_bar.update(1)\n            logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0], \"step\": global_step}\n            progress_bar.set_postfix(**logs)\n            accelerator.log(logs, step=global_step)\n            global_step += 1\n\n        # After each epoch you optionally sample some demo images with evaluate() and save the model\n        if accelerator.is_main_process:\n            pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)\n\n            if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:\n                evaluate(config, epoch, pipeline, vae)\n\n            if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:\n                if config.push_to_hub:\n                    pipeline.save_pretrained(config.output_dir)\n                    repo.push_to_hub(commit_message=f\"Epoch {epoch}\", blocking=True)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:02:09.924569Z","iopub.execute_input":"2023-09-24T16:02:09.924980Z","iopub.status.idle":"2023-09-24T16:02:09.950983Z","shell.execute_reply.started":"2023-09-24T16:02:09.924944Z","shell.execute_reply":"2023-09-24T16:02:09.949828Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def train_loop(config, model, vae, noise_scheduler, optimizer, train_dataloader, lr_scheduler):\n    # Initialize accelerator and tensorboard logging\n    accelerator = Accelerator(\n        mixed_precision=config.mixed_precision,\n        gradient_accumulation_steps=config.gradient_accumulation_steps,\n        log_with=\"tensorboard\",\n        project_dir=os.path.join(config.output_dir, \"logs\"),\n    )\n    if accelerator.is_main_process:\n        if config.push_to_hub:\n            repo_name = get_full_repo_name(Path(config.output_dir).name)\n            repo = Repository(config.output_dir, clone_from=repo_name)\n        elif config.output_dir is not None:\n            os.makedirs(config.output_dir, exist_ok=True)\n        accelerator.init_trackers(\"train_example\")\n\n    # Prepare everything\n    # There is no specific order to remember, you just need to unpack the\n    # objects in the same order you gave them to the prepare method.\n    model, vae, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n        model, vae, optimizer, train_dataloader, lr_scheduler\n    )\n    print(model.device)\n    global_step = 0\n\n    # Now you train the model\n    for epoch in range(config.num_epochs):\n        progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)\n        progress_bar.set_description(f\"Epoch {epoch}\")\n\n        for step, batch in enumerate(train_dataloader):\n            clean_images = batch[\"images\"]\n            \n            # Encode the images into latents [bs, 4, 32, 32]\n            latents = vae.encode(clean_images).latent_dist.sample()\n            \n            # Sample noise to add to the latents\n            noise = torch.randn(latents.shape).to(latents.device)\n            bs = latents.shape[0]\n\n            # Sample a random timestep for each image\n            timesteps = torch.randint(\n                0, \n                noise_scheduler.config.num_train_timesteps//2, # using less noisy images\n                (bs,), \n                device=clean_images.device\n            ).long()\n\n            # Add noise to the latens according to the noise magnitude at each timestep\n            # (this is the forward diffusion process)\n            noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)\n\n            with accelerator.accumulate(model):\n                # Predict the noise residual\n                noise_pred = model(noisy_latents, timesteps, return_dict=False)[0]\n                loss = F.mse_loss(noise_pred, noise)\n                accelerator.backward(loss)\n\n                accelerator.clip_grad_norm_(model.parameters(), 1.0)\n                optimizer.step()\n                lr_scheduler.step()\n                optimizer.zero_grad()\n\n            progress_bar.update(1)\n            logs = {\"loss\": loss.detach().item(), \"lr\": lr_scheduler.get_last_lr()[0], \"step\": global_step}\n            progress_bar.set_postfix(**logs)\n            accelerator.log(logs, step=global_step)\n            global_step += 1\n","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"config.train_batch_size = 2\nconfig.eval_batch_size = 2\nconfig.num_epochs=50\nsave_image_epochs = 10\nsave_model_epochs = 10","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:02:11.102064Z","iopub.execute_input":"2023-09-24T16:02:11.102463Z","iopub.status.idle":"2023-09-24T16:02:11.109710Z","shell.execute_reply.started":"2023-09-24T16:02:11.102431Z","shell.execute_reply":"2023-09-24T16:02:11.108514Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from accelerate import notebook_launcher\n\nargs = (config, model, vae, noise_scheduler, optimizer, train_dataloader, lr_scheduler)\n\nnotebook_launcher(train_loop, args, num_processes=1)","metadata":{"execution":{"iopub.status.busy":"2023-09-24T16:02:16.640510Z","iopub.execute_input":"2023-09-24T16:02:16.641744Z","iopub.status.idle":"2023-09-24T16:37:17.692736Z","shell.execute_reply.started":"2023-09-24T16:02:16.641706Z","shell.execute_reply":"2023-09-24T16:37:17.690175Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Once training is complete, take a look at the final monets generated by your diffusion model!","metadata":{}},{"cell_type":"code","source":"import glob\n\nsample_images = sorted(glob.glob(f\"{config.output_dir}/samples/*.png\"))\nImage.open(sample_images[-1])","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:43:09.996396Z","iopub.status.idle":"2023-09-24T15:43:09.997650Z","shell.execute_reply.started":"2023-09-24T15:43:09.997383Z","shell.execute_reply":"2023-09-24T15:43:09.997408Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"These are some pretty amazing monets conjured by our Stable diffusion model from random noise! We can clearly see some structures like water bodies and trees are starting to come come up by 100th epoch.","metadata":{}},{"cell_type":"markdown","source":"And if you like the notebook, please upvote ♥️","metadata":{}},{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"code","source":"#loading model from the hub to perform inference\nddpm = DDPMPipeline.from_pretrained(\n    \"NavneetSajwan/latents-trained-monets-256\",\n    use_safetensors=True\n)\n","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:43:10.001362Z","iopub.status.idle":"2023-09-24T15:43:10.002205Z","shell.execute_reply.started":"2023-09-24T15:43:10.001930Z","shell.execute_reply":"2023-09-24T15:43:10.001965Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"vae = AutoencoderKL.from_pretrained(\"CompVis/stable-diffusion-v1-1\",\n                                    subfolder = \"vae\",\n                                    safe_tensors=True\n                                   )","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:43:10.003655Z","iopub.status.idle":"2023-09-24T15:43:10.004472Z","shell.execute_reply.started":"2023-09-24T15:43:10.004212Z","shell.execute_reply":"2023-09-24T15:43:10.004238Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"images = ddpm(num_inference_steps=40, batch_size=1, output_type=\"array\").images\nlatents = torch.from_numpy(image.transpose(0,3, 1, 2))\ndecodedimgs = latents_to_pil(latents.to(vae.device), vae)\ndecodedimgs[0]","metadata":{"execution":{"iopub.status.busy":"2023-09-24T15:43:10.005959Z","iopub.status.idle":"2023-09-24T15:43:10.006737Z","shell.execute_reply.started":"2023-09-24T15:43:10.006477Z","shell.execute_reply":"2023-09-24T15:43:10.006501Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}