%pip install -qq diffusers datasets accelerate wandb open-clip-torch

>>> # Code to log in to the Hugging Face Hub, needed for sharing models
>>> # Make sure you use a token with WRITE access
>>> from huggingface_hub import notebook_login

>>> notebook_login()

import numpy as np
import torch
import torch.nn.functional as F
import torchvision
from datasets import load_dataset
from diffusers import DDIMScheduler, DDPMPipeline
from matplotlib import pyplot as plt
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm

device = "mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"

image_pipe = DDPMPipeline.from_pretrained("google/ddpm-celebahq-256")
image_pipe.to(device)

>>> images = image_pipe().images
>>> images[0]

# Create new scheduler and set num inference steps
scheduler = DDIMScheduler.from_pretrained("google/ddpm-celebahq-256")
scheduler.set_timesteps(num_inference_steps=40)

scheduler.timesteps

>>> # The random starting point
>>> x = torch.randn(4, 3, 256, 256).to(device)  # Batch of 4, 3-channel 256 x 256 px images

>>> # Loop through the sampling timesteps
>>> for i, t in tqdm(enumerate(scheduler.timesteps)):

...     # Prepare model input
...     model_input = scheduler.scale_model_input(x, t)

...     # Get the prediction
...     with torch.no_grad():
...         noise_pred = image_pipe.unet(model_input, t)["sample"]

...     # Calculate what the updated sample should look like with the scheduler
...     scheduler_output = scheduler.step(noise_pred, t, x)

...     # Update x
...     x = scheduler_output.prev_sample

...     # Occasionally display both x and the predicted denoised images
...     if i % 10 == 0 or i == len(scheduler.timesteps) - 1:
...         fig, axs = plt.subplots(1, 2, figsize=(12, 5))

...         grid = torchvision.utils.make_grid(x, nrow=4).permute(1, 2, 0)
...         axs[0].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
...         axs[0].set_title(f"Current x (step {i})")

...         pred_x0 = scheduler_output.pred_original_sample  # Not available for all schedulers
...         grid = torchvision.utils.make_grid(pred_x0, nrow=4).permute(1, 2, 0)
...         axs[1].imshow(grid.cpu().clip(-1, 1) * 0.5 + 0.5)
...         axs[1].set_title(f"Predicted denoised images (step {i})")
...         plt.show()

# ??scheduler.step

>>> image_pipe.scheduler = scheduler
>>> images = image_pipe(num_inference_steps=40).images
>>> images[0]

>>> # @markdown load and prepare a dataset:
>>> # Not on Colab? Comments with #@ enable UI tweaks like headings or user inputs
>>> # but can safely be ignored if you're working on a different platform.

>>> dataset_name = "huggan/smithsonian_butterflies_subset"  # @param
>>> dataset = load_dataset(dataset_name, split="train")
>>> image_size = 256  # @param
>>> batch_size = 4  # @param
>>> preprocess = transforms.Compose(
...     [
...         transforms.Resize((image_size, image_size)),
...         transforms.RandomHorizontalFlip(),
...         transforms.ToTensor(),
...         transforms.Normalize([0.5], [0.5]),
...     ]
... )


>>> def transform(examples):
...     images = [preprocess(image.convert("RGB")) for image in examples["image"]]
...     return {"images": images}


>>> dataset.set_transform(transform)

>>> train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

>>> print("Previewing batch:")
>>> batch = next(iter(train_dataloader))
>>> grid = torchvision.utils.make_grid(batch["images"], nrow=4)
>>> plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5)

>>> num_epochs = 2  # @param
>>> lr = 1e-5  # 2param
>>> grad_accumulation_steps = 2  # @param

>>> optimizer = torch.optim.AdamW(image_pipe.unet.parameters(), lr=lr)

>>> losses = []

>>> for epoch in range(num_epochs):
...     for step, batch in tqdm(enumerate(train_dataloader), total=len(train_dataloader)):
...         clean_images = batch["images"].to(device)
...         # Sample noise to add to the images
...         noise = torch.randn(clean_images.shape).to(clean_images.device)
...         bs = clean_images.shape[0]

...         # Sample a random timestep for each image
...         timesteps = torch.randint(
...             0,
...             image_pipe.scheduler.num_train_timesteps,
...             (bs,),
...             device=clean_images.device,
...         ).long()

...         # Add noise to the clean images according to the noise magnitude at each timestep
...         # (this is the forward diffusion process)
...         noisy_images = image_pipe.scheduler.add_noise(clean_images, noise, timesteps)

...         # Get the model prediction for the noise
...         noise_pred = image_pipe.unet(noisy_images, timesteps, return_dict=False)[0]

...         # Compare the prediction with the actual noise:
...         loss = F.mse_loss(
...             noise_pred, noise
...         )  # NB - trying to predict noise (eps) not (noisy_ims-clean_ims) or just (clean_ims)

...         # Store for later plotting
...         losses.append(loss.item())

...         # Update the model parameters with the optimizer based on this loss
...         loss.backward(loss)

...         # Gradient accumulation:
...         if (step + 1) % grad_accumulation_steps == 0:
...             optimizer.step()
...             optimizer.zero_grad()

...     print(f"Epoch {epoch} average loss: {sum(losses[-len(train_dataloader):])/len(train_dataloader)}")

>>> # Plot the loss curve:
>>> plt.plot(losses)

# Exercise: See if you can add gradient accumulation to the training loop in Unit 1.
# How does it perform? Think how you might adjust the learning rate based on the
# number of gradient accumulation steps - should it stay the same as before?

%wandb johnowhitaker/dm_finetune/2upaa341 # You'll need a W&B account for this to work - skip if you don't want to log in

# Exercise: see if you can modify the official example training script we saw
# in Unit 1 to begin with a pre-trained model rather than training from scratch.
# Compare it to the minimal script linked above - what extra features is the minimal script missing?

>>> # @markdown Generate and plot some images:
>>> x = torch.randn(8, 3, 256, 256).to(device)  # Batch of 8
>>> for i, t in tqdm(enumerate(scheduler.timesteps)):
...     model_input = scheduler.scale_model_input(x, t)
...     with torch.no_grad():
...         noise_pred = image_pipe.unet(model_input, t)["sample"]
...     x = scheduler.step(noise_pred, t, x).prev_sample
>>> grid = torchvision.utils.make_grid(x, nrow=4)
>>> plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5)

## To download the fine-tuning script:
# !wget https://github.com/huggingface/diffusion-models-class/raw/main/unit2/finetune_model.py

## To run the script, training the face model on some vintage faces
## (ideally run this in a terminal):
# !python finetune_model.py --image_size 128 --batch_size 8 --num_epochs 16\
#     --grad_accumulation_steps 2 --start_model "google/ddpm-celebahq-256"\
#     --dataset_name "Norod78/Vintage-Faces-FFHQAligned" --wandb_project 'dm-finetune'\
#     --log_samples_every 100 --save_model_every 1000 --model_save_name 'vintageface'

image_pipe.save_pretrained("my-finetuned-model")

>>> !ls {"my-finetuned-model"}

# @title Upload a locally saved pipeline to the hub

# Code to upload a pipeline saved locally to the hub
from huggingface_hub import HfApi, ModelCard, create_repo, get_full_repo_name

# Set up repo and upload files
model_name = "ddpm-celebahq-finetuned-butterflies-2epochs"  # @param What you want it called on the hub
local_folder_name = (
    "my-finetuned-model"  # @param Created by the script or one you created via image_pipe.save_pretrained('save_name')
)
description = "Describe your model here"  # @param
hub_model_id = get_full_repo_name(model_name)
create_repo(hub_model_id)
api = HfApi()
api.upload_folder(folder_path=f"{local_folder_name}/scheduler", path_in_repo="", repo_id=hub_model_id)
api.upload_folder(folder_path=f"{local_folder_name}/unet", path_in_repo="", repo_id=hub_model_id)
api.upload_file(
    path_or_fileobj=f"{local_folder_name}/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

# Add a model card (optional but nice!)
content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

# Example Fine-Tuned Model for Unit 2 of the [Diffusion Models Class 🧨](https://github.com/huggingface/diffusion-models-class)

{description}

## Usage

```python
from diffusers import DDPMPipeline

pipeline = DDPMPipeline.from_pretrained('{hub_model_id}')
image = pipeline().images[0]
image

Congratulations, you've now fine-tuned your first diffusion model!

For the rest of this notebook we'll use a [model](https://huggingface.co/johnowhitaker/sd-class-wikiart-from-bedrooms) I fine-tuned from [this model trained on LSUN bedrooms](https://huggingface.co/google/ddpm-bedroom-256) approximately one epoch on the [WikiArt dataset](https://huggingface.co/datasets/huggan/wikiart). If you'd prefer, you can skip this cell and use the faces/butterflies pipeline we fine-tuned in the previous section or load one from the Hub instead:

```python
>>> # Load the pretrained pipeline
>>> pipeline_name = "johnowhitaker/sd-class-wikiart-from-bedrooms"
>>> image_pipe = DDPMPipeline.from_pretrained(pipeline_name).to(device)

>>> # Sample some images with a DDIM Scheduler over 40 steps
>>> scheduler = DDIMScheduler.from_pretrained(pipeline_name)
>>> scheduler.set_timesteps(num_inference_steps=40)

>>> # Random starting point (batch of 8 images)
>>> x = torch.randn(8, 3, 256, 256).to(device)

>>> # Minimal sampling loop
>>> for i, t in tqdm(enumerate(scheduler.timesteps)):
...     model_input = scheduler.scale_model_input(x, t)
...     with torch.no_grad():
...         noise_pred = image_pipe.unet(model_input, t)["sample"]
...     x = scheduler.step(noise_pred, t, x).prev_sample

>>> # View the results
>>> grid = torchvision.utils.make_grid(x, nrow=4)
>>> plt.imshow(grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5)

def color_loss(images, target_color=(0.1, 0.9, 0.5)):
    """Given a target color (R, G, B) return a loss for how far away on average
    the images' pixels are from that color. Defaults to a light teal: (0.1, 0.9, 0.5)"""
    target = torch.tensor(target_color).to(images.device) * 2 - 1  # Map target color to (-1, 1)
    target = target[None, :, None, None]  # Get shape right to work with the images (b, c, h, w)
    error = torch.abs(images - target).mean()  # Mean absolute difference between the image pixels and the target color
    return error

>>> # Variant 1: shortcut method

>>> # The guidance scale determines the strength of the effect
>>> guidance_loss_scale = 40  # Explore changing this to 5, or 100

>>> x = torch.randn(8, 3, 256, 256).to(device)

>>> for i, t in tqdm(enumerate(scheduler.timesteps)):

...     # Prepare the model input
...     model_input = scheduler.scale_model_input(x, t)

...     # predict the noise residual
...     with torch.no_grad():
...         noise_pred = image_pipe.unet(model_input, t)["sample"]

...     # Set x.requires_grad to True
...     x = x.detach().requires_grad_()

...     # Get the predicted x0
...     x0 = scheduler.step(noise_pred, t, x).pred_original_sample

...     # Calculate loss
...     loss = color_loss(x0) * guidance_loss_scale
...     if i % 10 == 0:
...         print(i, "loss:", loss.item())

...     # Get gradient
...     cond_grad = -torch.autograd.grad(loss, x)[0]

...     # Modify x based on this gradient
...     x = x.detach() + cond_grad

...     # Now step with scheduler
...     x = scheduler.step(noise_pred, t, x).prev_sample

>>> # View the output
>>> grid = torchvision.utils.make_grid(x, nrow=4)
>>> im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
>>> Image.fromarray(np.array(im * 255).astype(np.uint8))

>>> # Variant 2: setting x.requires_grad before calculating the model predictions

>>> guidance_loss_scale = 40
>>> x = torch.randn(4, 3, 256, 256).to(device)

>>> for i, t in tqdm(enumerate(scheduler.timesteps)):

...     # Set requires_grad before the model forward pass
...     x = x.detach().requires_grad_()
...     model_input = scheduler.scale_model_input(x, t)

...     # predict (with grad this time)
...     noise_pred = image_pipe.unet(model_input, t)["sample"]

...     # Get the predicted x0:
...     x0 = scheduler.step(noise_pred, t, x).pred_original_sample

...     # Calculate loss
...     loss = color_loss(x0) * guidance_loss_scale
...     if i % 10 == 0:
...         print(i, "loss:", loss.item())

...     # Get gradient
...     cond_grad = -torch.autograd.grad(loss, x)[0]

...     # Modify x based on this gradient
...     x = x.detach() + cond_grad

...     # Now step with scheduler
...     x = scheduler.step(noise_pred, t, x).prev_sample


>>> grid = torchvision.utils.make_grid(x, nrow=4)
>>> im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
>>> Image.fromarray(np.array(im * 255).astype(np.uint8))

# Exercise: pick your favourite colour and look up it's values in RGB space.
# Edit the `color_loss()` line in the cell above to receive these new RGB values and examine the outputs - do they match what you expect?

# @markdown load a CLIP model and define the loss function
import open_clip

clip_model, _, preprocess = open_clip.create_model_and_transforms("ViT-B-32", pretrained="openai")
clip_model.to(device)

# Transforms to resize and augment an image + normalize to match CLIP's training data
tfms = torchvision.transforms.Compose(
    [
        torchvision.transforms.RandomResizedCrop(224),  # Random CROP each time
        torchvision.transforms.RandomAffine(5),  # One possible random augmentation: skews the image
        torchvision.transforms.RandomHorizontalFlip(),  # You can add additional augmentations if you like
        torchvision.transforms.Normalize(
            mean=(0.48145466, 0.4578275, 0.40821073),
            std=(0.26862954, 0.26130258, 0.27577711),
        ),
    ]
)


# And define a loss function that takes an image, embeds it and compares with
# the text features of the prompt
def clip_loss(image, text_features):
    image_features = clip_model.encode_image(tfms(image))  # Note: applies the above transforms
    input_normed = torch.nn.functional.normalize(image_features.unsqueeze(1), dim=2)
    embed_normed = torch.nn.functional.normalize(text_features.unsqueeze(0), dim=2)
    dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)  # Squared Great Circle Distance
    return dists.mean()

>>> # @markdown applying guidance using CLIP

>>> prompt = "Red Rose (still life), red flower painting"  # @param

>>> # Explore changing this
>>> guidance_scale = 8  # @param
>>> n_cuts = 4  # @param

>>> # More steps -> more time for the guidance to have an effect
>>> scheduler.set_timesteps(50)

>>> # We embed a prompt with CLIP as our target
>>> text = open_clip.tokenize([prompt]).to(device)
>>> with torch.no_grad(), torch.cuda.amp.autocast():
...     text_features = clip_model.encode_text(text)


>>> x = torch.randn(4, 3, 256, 256).to(device)  # RAM usage is high, you may want only 1 image at a time

>>> for i, t in tqdm(enumerate(scheduler.timesteps)):

...     model_input = scheduler.scale_model_input(x, t)

...     # predict the noise residual
...     with torch.no_grad():
...         noise_pred = image_pipe.unet(model_input, t)["sample"]

...     cond_grad = 0

...     for cut in range(n_cuts):

...         # Set requires grad on x
...         x = x.detach().requires_grad_()

...         # Get the predicted x0:
...         x0 = scheduler.step(noise_pred, t, x).pred_original_sample

...         # Calculate loss
...         loss = clip_loss(x0, text_features) * guidance_scale

...         # Get gradient (scale by n_cuts since we want the average)
...         cond_grad -= torch.autograd.grad(loss, x)[0] / n_cuts

...     if i % 25 == 0:
...         print("Step:", i, ", Guidance loss:", loss.item())

...     # Modify x based on this gradient
...     alpha_bar = scheduler.alphas_cumprod[i]
...     x = x.detach() + cond_grad * alpha_bar.sqrt()  # Note the additional scaling factor here!

...     # Now step with scheduler
...     x = scheduler.step(noise_pred, t, x).prev_sample


>>> grid = torchvision.utils.make_grid(x.detach(), nrow=4)
>>> im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
>>> Image.fromarray(np.array(im * 255).astype(np.uint8))

>>> # @markdown Plotting some possible schedules:
>>> plt.plot([1 for a in scheduler.alphas_cumprod], label="no scaling")
>>> plt.plot([a for a in scheduler.alphas_cumprod], label="alpha_bar")
>>> plt.plot([a.sqrt() for a in scheduler.alphas_cumprod], label="alpha_bar.sqrt()")
>>> plt.plot([(1 - a).sqrt() for a in scheduler.alphas_cumprod], label="(1-alpha_bar).sqrt()")
>>> plt.legend()
>>> plt.title("Possible guidance scaling schedules")

%pip install -q gradio # Install the library

import gradio as gr
from PIL import Image, ImageColor


# The function that does the hard work
def generate(color, guidance_loss_scale):
    target_color = ImageColor.getcolor(color, "RGB")  # Target color as RGB
    target_color = [a / 255 for a in target_color]  # Rescale from (0, 255) to (0, 1)
    x = torch.randn(1, 3, 256, 256).to(device)
    for i, t in tqdm(enumerate(scheduler.timesteps)):
        model_input = scheduler.scale_model_input(x, t)
        with torch.no_grad():
            noise_pred = image_pipe.unet(model_input, t)["sample"]
        x = x.detach().requires_grad_()
        x0 = scheduler.step(noise_pred, t, x).pred_original_sample
        loss = color_loss(x0, target_color) * guidance_loss_scale
        cond_grad = -torch.autograd.grad(loss, x)[0]
        x = x.detach() + cond_grad
        x = scheduler.step(noise_pred, t, x).prev_sample
    grid = torchvision.utils.make_grid(x, nrow=4)
    im = grid.permute(1, 2, 0).cpu().clip(-1, 1) * 0.5 + 0.5
    im = Image.fromarray(np.array(im * 255).astype(np.uint8))
    im.save("test.jpeg")
    return im


# See the gradio docs for the types of inputs and outputs available
inputs = [
    gr.ColorPicker(label="color", value="55FFAA"),  # Add any inputs you need here
    gr.Slider(label="guidance_scale", minimum=0, maximum=30, value=3),
]
outputs = gr.Image(label="result")

# And the minimal interface
demo = gr.Interface(
    fn=generate,
    inputs=inputs,
    outputs=outputs,
    examples=[
        ["#BB2266", 3],
        ["#44CCAA", 5],  # You can provide some example inputs to get people started
    ],
)
demo.launch(debug=True)  # debug=True allows you to see errors and output in Colab