%pip install -qq -U diffusers datasets transformers accelerate ftfy pyarrow==9.0.0

>>> from huggingface_hub import notebook_login

>>> notebook_login()

%%capture
!sudo apt -qq install git-lfs
!git config --global credential.helper store

import numpy as np
import torch
import torch.nn.functional as F
from matplotlib import pyplot as plt
from PIL import Image


def show_images(x):
    """Given a batch of images x, make a grid and convert to PIL"""
    x = x * 0.5 + 0.5  # Map from (-1, 1) back to (0, 1)
    grid = torchvision.utils.make_grid(x)
    grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
    grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
    return grid_im


def make_grid(images, size=64):
    """Given a list of PIL images, stack them together into a line for easy viewing"""
    output_im = Image.new("RGB", (size * len(images), size))
    for i, im in enumerate(images):
        output_im.paste(im.resize((size, size)), (i * size, 0))
    return output_im


# Mac users may need device = 'mps' (untested)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

>>> from IPython.display import YouTubeVideo

>>> YouTubeVideo("W4Mcuh38wyM")

from diffusers import StableDiffusionPipeline

# Check out https://huggingface.co/sd-dreambooth-library for loads of models from the community
model_id = "sd-dreambooth-library/mr-potato-head"

# Load the pipeline
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16).to(device)

>>> prompt = "an abstract oil painting of sks mr potato head by picasso"
>>> image = pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]
>>> image

>>> from diffusers import DDPMPipeline

>>> # Load the butterfly pipeline
>>> butterfly_pipeline = DDPMPipeline.from_pretrained("johnowhitaker/ddpm-butterflies-32px").to(device)

>>> # Create 8 images
>>> images = butterfly_pipeline(batch_size=8).images

>>> # View the result
>>> make_grid(images)

import torchvision
from datasets import load_dataset
from torchvision import transforms

dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")

# Or load images from a local folder
# dataset = load_dataset("imagefolder", data_dir="path/to/folder")

# We'll train on 32-pixel square images, but you can try larger sizes too
image_size = 32
# You can lower your batch size if you're running out of GPU memory
batch_size = 64

# Define data augmentations
preprocess = transforms.Compose(
    [
        transforms.Resize((image_size, image_size)),  # Resize
        transforms.RandomHorizontalFlip(),  # Randomly flip (data augmentation)
        transforms.ToTensor(),  # Convert to tensor (0, 1)
        transforms.Normalize([0.5], [0.5]),  # Map to (-1, 1)
    ]
)


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


dataset.set_transform(transform)

# Create a dataloader from the dataset to serve up the transformed images in batches
train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

>>> xb = next(iter(train_dataloader))["images"].to(device)[:8]
>>> print("X shape:", xb.shape)
>>> show_images(xb).resize((8 * 64, 64), resample=Image.NEAREST)

from diffusers import DDPMScheduler

noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

>>> plt.plot(noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r"${\sqrt{\bar{\alpha}_t}}$")
>>> plt.plot((1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5, label=r"$\sqrt{(1 - \bar{\alpha}_t)}$")
>>> plt.legend(fontsize="x-large")

# One with too little noise added:
# noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_start=0.001, beta_end=0.004)
# The 'cosine' schedule, which may be better for small image sizes:
# noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule='squaredcos_cap_v2')

>>> timesteps = torch.linspace(0, 999, 8).long().to(device)
>>> noise = torch.randn_like(xb)
>>> noisy_xb = noise_scheduler.add_noise(xb, noise, timesteps)
>>> print("Noisy X shape", noisy_xb.shape)
>>> show_images(noisy_xb).resize((8 * 64, 64), resample=Image.NEAREST)

from diffusers import UNet2DModel

# Create a model
model = UNet2DModel(
    sample_size=image_size,  # the target image resolution
    in_channels=3,  # the number of input channels, 3 for RGB images
    out_channels=3,  # the number of output channels
    layers_per_block=2,  # how many ResNet layers to use per UNet block
    block_out_channels=(64, 128, 128, 256),  # More channels -> more parameters
    down_block_types=(
        "DownBlock2D",  # a regular ResNet downsampling block
        "DownBlock2D",
        "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
        "AttnDownBlock2D",
    ),
    up_block_types=(
        "AttnUpBlock2D",
        "AttnUpBlock2D",  # a ResNet upsampling block with spatial self-attention
        "UpBlock2D",
        "UpBlock2D",  # a regular ResNet upsampling block
    ),
)
model.to(device)

with torch.no_grad():
    model_prediction = model(noisy_xb, timesteps).sample
model_prediction.shape

>>> # Set the noise scheduler
>>> noise_scheduler = DDPMScheduler(num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2")

>>> # Training loop
>>> optimizer = torch.optim.AdamW(model.parameters(), lr=4e-4)

>>> losses = []

>>> for epoch in range(30):
...     for step, batch in enumerate(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, noise_scheduler.num_train_timesteps, (bs,), device=clean_images.device).long()

...         # Add noise to the clean images according to the noise magnitude at each timestep
...         noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

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

...         # Calculate the loss
...         loss = F.mse_loss(noise_pred, noise)
...         loss.backward(loss)
...         losses.append(loss.item())

...         # Update the model parameters with the optimizer
...         optimizer.step()
...         optimizer.zero_grad()

...     if (epoch + 1) % 5 == 0:
...         loss_last_epoch = sum(losses[-len(train_dataloader) :]) / len(train_dataloader)
...         print(f"Epoch:{epoch+1}, loss: {loss_last_epoch}")

>>> fig, axs = plt.subplots(1, 2, figsize=(12, 4))
>>> axs[0].plot(losses)
>>> axs[1].plot(np.log(losses))
>>> plt.show()

# Uncomment to instead load the model I trained earlier:
# model = butterfly_pipeline.unet

from diffusers import DDPMPipeline

image_pipe = DDPMPipeline(unet=model, scheduler=noise_scheduler)

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

image_pipe.save_pretrained("my_pipeline")

>>> !ls my_pipeline/

>>> !ls my_pipeline/unet/

# ??image_pipe.forward

>>> # Random starting point (8 random images):
>>> sample = torch.randn(8, 3, 32, 32).to(device)

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

...     # Get model pred
...     with torch.no_grad():
...         residual = model(sample, t).sample

...     # Update sample with step
...     sample = noise_scheduler.step(residual, t, sample).prev_sample

>>> show_images(sample)

from huggingface_hub import get_full_repo_name

model_name = "sd-class-butterflies-32"
hub_model_id = get_full_repo_name(model_name)
hub_model_id

from huggingface_hub import HfApi, create_repo

create_repo(hub_model_id)
api = HfApi()
api.upload_folder(folder_path="my_pipeline/scheduler", path_in_repo="", repo_id=hub_model_id)
api.upload_folder(folder_path="my_pipeline/unet", path_in_repo="", repo_id=hub_model_id)
api.upload_file(
    path_or_fileobj="my_pipeline/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

from huggingface_hub import ModelCard

content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

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

This model is a diffusion model for unconditional image generation of cute 🦋.

## Usage

```python
from diffusers import DDPMPipeline

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

Now that the model is on the Hub, you can download it from anywhere by using the `from_pretrained()` method of the `DDPMPipeline` as follows"

```python
>>> from diffusers import DDPMPipeline

>>> image_pipe = DDPMPipeline.from_pretrained(hub_model_id)
>>> pipeline_output = image_pipe()
>>> pipeline_output.images[0]

!wget https://github.com/huggingface/diffusers/raw/main/examples/unconditional_image_generation/train_unconditional.py

# Let's give our new model a name for the Hub
model_name = "sd-class-butterflies-64"
hub_model_id = get_full_repo_name(model_name)
hub_model_id

!accelerate launch train_unconditional.py \
  --dataset_name="huggan/smithsonian_butterflies_subset" \
  --resolution=64 \
  --output_dir={model_name} \
  --train_batch_size=32 \
  --num_epochs=50 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision="no"

create_repo(hub_model_id)
api = HfApi()
api.upload_folder(folder_path=f"{model_name}/scheduler", path_in_repo="", repo_id=hub_model_id)
api.upload_folder(folder_path=f"{model_name}/unet", path_in_repo="", repo_id=hub_model_id)
api.upload_file(
    path_or_fileobj=f"{model_name}/model_index.json",
    path_in_repo="model_index.json",
    repo_id=hub_model_id,
)

content = f"""
---
license: mit
tags:
- pytorch
- diffusers
- unconditional-image-generation
- diffusion-models-class
---

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

This model is a diffusion model for unconditional image generation of cute 🦋.

## Usage

```python
from diffusers import DDPMPipeline

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

About 45 minutes later, this is the result:

```python
>>> pipeline = DDPMPipeline.from_pretrained(hub_model_id).to(device)
>>> images = pipeline(batch_size=8).images
>>> make_grid(images)