{"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":"## 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-18T06:30:41.339671Z","iopub.execute_input":"2023-09-18T06:30:41.340016Z","iopub.status.idle":"2023-09-18T06:31:42.065079Z","shell.execute_reply.started":"2023-09-18T06:30:41.339986Z","shell.execute_reply":"2023-09-18T06:31:42.063852Z"},"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-18T06:31:42.067924Z","iopub.execute_input":"2023-09-18T06:31:42.069399Z","iopub.status.idle":"2023-09-18T06:31:42.432920Z","shell.execute_reply.started":"2023-09-18T06:31:42.069367Z","shell.execute_reply":"2023-09-18T06:31:42.431982Z"},"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-18T06:31:57.199124Z","iopub.execute_input":"2023-09-18T06:31:57.199826Z","iopub.status.idle":"2023-09-18T06:32:01.360032Z","shell.execute_reply.started":"2023-09-18T06:31:57.199788Z","shell.execute_reply":"2023-09-18T06:32:01.358419Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Preprocessing","metadata":{}},{"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-18T06:32:41.554066Z","iopub.execute_input":"2023-09-18T06:32:41.554476Z","iopub.status.idle":"2023-09-18T06:32:46.380487Z","shell.execute_reply.started":"2023-09-18T06:32:41.554444Z","shell.execute_reply":"2023-09-18T06:32:46.379301Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#let's inspect our dataset\ndataset, len(dataset)","metadata":{"execution":{"iopub.status.busy":"2023-09-18T06:34:29.264392Z","iopub.execute_input":"2023-09-18T06:34:29.264794Z","iopub.status.idle":"2023-09-18T06:34:29.274334Z","shell.execute_reply.started":"2023-09-18T06:34:29.264760Z","shell.execute_reply":"2023-09-18T06:34:29.273396Z"},"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":"code","source":"#let's inspect a single image by indexing into our dataset\ndataset[0]","metadata":{"execution":{"iopub.status.busy":"2023-09-18T06:35:06.338178Z","iopub.execute_input":"2023-09-18T06:35:06.339263Z","iopub.status.idle":"2023-09-18T06:35:06.351148Z","shell.execute_reply.started":"2023-09-18T06:35:06.339191Z","shell.execute_reply":"2023-09-18T06:35:06.349910Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":"markdown","source":"## Stage 1 training: 128 * 128 input","metadata":{}},{"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 = 128  # the generated image resolution\n    train_batch_size = 16\n    eval_batch_size = 16  # how many images to sample during evaluation\n    num_epochs = 100\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 = \"unconditional-monet-128\"  # 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-18T06:39:19.636630Z","iopub.execute_input":"2023-09-18T06:39:19.637030Z","iopub.status.idle":"2023-09-18T06:39:19.644368Z","shell.execute_reply.started":"2023-09-18T06:39:19.637000Z","shell.execute_reply":"2023-09-18T06:39:19.643300Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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-18T06:39:23.923188Z","iopub.execute_input":"2023-09-18T06:39:23.924056Z","iopub.status.idle":"2023-09-18T06:39:24.499815Z","shell.execute_reply.started":"2023-09-18T06:39:23.924010Z","shell.execute_reply":"2023-09-18T06:39:24.498827Z"},"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((config.image_size, config.image_size)),\n        transforms.RandomHorizontalFlip(),\n        transforms.ToTensor(),\n        transforms.Normalize([0.5], [0.5]),\n    ]\n)","metadata":{"execution":{"iopub.status.busy":"2023-09-18T06:39:29.419333Z","iopub.execute_input":"2023-09-18T06:39:29.419711Z","iopub.status.idle":"2023-09-18T06:39:35.095281Z","shell.execute_reply.started":"2023-09-18T06:39:29.419681Z","shell.execute_reply":"2023-09-18T06:39:35.094292Z"},"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-18T06:39:37.261307Z","iopub.execute_input":"2023-09-18T06:39:37.261897Z","iopub.status.idle":"2023-09-18T06:39:37.268903Z","shell.execute_reply.started":"2023-09-18T06:39:37.261862Z","shell.execute_reply":"2023-09-18T06:39:37.267709Z"},"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-18T06:39:37.419917Z","iopub.execute_input":"2023-09-18T06:39:37.420277Z","iopub.status.idle":"2023-09-18T06:39:37.425751Z","shell.execute_reply.started":"2023-09-18T06:39:37.420226Z","shell.execute_reply":"2023-09-18T06:39:37.424565Z"},"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":"from diffusers import UNet2DModel\n\nmodel = UNet2DModel(\n    sample_size=config.image_size,  # the target image resolution\n    in_channels=3,  # the number of input channels, 3 for RGB images\n    out_channels=3,  # 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-18T06:39:41.475304Z","iopub.execute_input":"2023-09-18T06:39:41.475666Z","iopub.status.idle":"2023-09-18T06:39:57.776662Z","shell.execute_reply.started":"2023-09-18T06:39:41.475635Z","shell.execute_reply":"2023-09-18T06:39:57.775329Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# check if the input and and output are of same size\nsample_image = dataset[0][\"images\"].unsqueeze(0)\nprint(\"Input shape:\", sample_image.shape)\n\nprint(\"Output shape:\", model(sample_image, timestep=0).sample.shape)","metadata":{"execution":{"iopub.status.busy":"2023-09-18T06:40:07.928450Z","iopub.execute_input":"2023-09-18T06:40:07.929433Z","iopub.status.idle":"2023-09-18T06:40:10.422369Z","shell.execute_reply.started":"2023-09-18T06:40:07.929379Z","shell.execute_reply":"2023-09-18T06:40:10.421308Z"},"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":"import torch\nfrom PIL import Image\nfrom diffusers import DDPMScheduler\n\nnoise_scheduler = DDPMScheduler(num_train_timesteps=1000)\nnoise = torch.randn(sample_image.shape)\ntimesteps = torch.LongTensor([50])\nnoisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)\n\nImage.fromarray(((noisy_image.permute(0, 2, 3, 1) + 1.0) * 127.5).type(torch.uint8).numpy()[0])","metadata":{"execution":{"iopub.status.busy":"2023-09-18T06:40:52.973938Z","iopub.execute_input":"2023-09-18T06:40:52.974341Z","iopub.status.idle":"2023-09-18T06:40:53.012554Z","shell.execute_reply.started":"2023-09-18T06:40:52.974307Z","shell.execute_reply":"2023-09-18T06:40:53.011467Z"},"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_image, timesteps).sample\nloss = F.mse_loss(noise_pred, noise)","metadata":{"execution":{"iopub.status.busy":"2023-09-18T06:40:57.148664Z","iopub.execute_input":"2023-09-18T06:40:57.149102Z","iopub.status.idle":"2023-09-18T06:40:59.641762Z","shell.execute_reply.started":"2023-09-18T06:40:57.149067Z","shell.execute_reply":"2023-09-18T06:40:59.640592Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train the model","metadata":{}},{"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-18T06:41:03.125910Z","iopub.execute_input":"2023-09-18T06:41:03.126300Z","iopub.status.idle":"2023-09-18T06:41:03.136925Z","shell.execute_reply.started":"2023-09-18T06:41:03.126262Z","shell.execute_reply":"2023-09-18T06:41:03.135654Z"},"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):\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        batch_size=config.eval_batch_size,\n        generator=torch.manual_seed(config.seed),\n    ).images\n\n    # Make a grid out of the images\n    image_grid = make_image_grid(images, rows=4, cols=4)\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-18T06:41:20.720345Z","iopub.execute_input":"2023-09-18T06:41:20.720938Z","iopub.status.idle":"2023-09-18T06:41:20.746967Z","shell.execute_reply.started":"2023-09-18T06:41:20.720907Z","shell.execute_reply":"2023-09-18T06:41:20.745926Z"},"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, 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, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(\n        model, optimizer, train_dataloader, lr_scheduler\n    )\n\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            # 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, noise_scheduler.config.num_train_timesteps, (bs,), 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)\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-18T06:42:12.454669Z","iopub.execute_input":"2023-09-18T06:42:12.455046Z","iopub.status.idle":"2023-09-18T06:42:12.474032Z","shell.execute_reply.started":"2023-09-18T06:42:12.455014Z","shell.execute_reply":"2023-09-18T06:42:12.472941Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from accelerate import notebook_launcher\n\nargs = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)\n\nnotebook_launcher(train_loop, args, num_processes=2)","metadata":{"execution":{"iopub.status.busy":"2023-09-18T06:42:24.471187Z","iopub.execute_input":"2023-09-18T06:42:24.471936Z","iopub.status.idle":"2023-09-18T07:47:51.515573Z","shell.execute_reply.started":"2023-09-18T06:42:24.471901Z","shell.execute_reply":"2023-09-18T07:47:51.514063Z"},"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-18T08:22:08.247786Z","iopub.execute_input":"2023-09-18T08:22:08.248270Z","iopub.status.idle":"2023-09-18T08:22:08.339903Z","shell.execute_reply.started":"2023-09-18T08:22:08.248210Z","shell.execute_reply":"2023-09-18T08:22:08.339066Z"},"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":"code","source":"#loading model from the hub to perform inference\nddpm = DDPMPipeline.from_pretrained(\"NavneetSajwan/unconditional-monet-128\", use_safetensors=True).to(\"cuda\")\nimage = ddpm(num_inference_steps=25).images[0]\nimage","metadata":{"execution":{"iopub.status.busy":"2023-09-18T08:22:18.543726Z","iopub.execute_input":"2023-09-18T08:22:18.544115Z","iopub.status.idle":"2023-09-18T08:22:35.632191Z","shell.execute_reply.started":"2023-09-18T08:22:18.544083Z","shell.execute_reply":"2023-09-18T08:22:35.631217Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image = ddpm(num_inference_steps=200 ).images[0]\nimage.resize((256,256))","metadata":{"execution":{"iopub.status.busy":"2023-09-18T08:24:49.787277Z","iopub.execute_input":"2023-09-18T08:24:49.787637Z","iopub.status.idle":"2023-09-18T08:24:57.863138Z","shell.execute_reply.started":"2023-09-18T08:24:49.787606Z","shell.execute_reply":"2023-09-18T08:24:57.862136Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}