%pip install -Uq diffusers ftfy accelerate

# Installing transformers from source for now since we need the latest version for Depth2Img
%pip install -Uq git+https://github.com/huggingface/transformers

import torch
import requests
from PIL import Image
from io import BytesIO
from matplotlib import pyplot as plt

# We'll be exploring a number of pipelines today!
from diffusers import (
    StableDiffusionPipeline,
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
    StableDiffusionDepth2ImgPipeline,
)


# We'll use a couple of demo images later in the notebook
def download_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert("RGB")


# Download images for inpainting example
img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = download_image(img_url).resize((512, 512))
mask_image = download_image(mask_url).resize((512, 512))

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

# Load the pipeline
model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)

>>> # Set up a generator for reproducibility
>>> generator = torch.Generator(device=device).manual_seed(42)

>>> # Run the pipeline, showing some of the available arguments
>>> pipe_output = pipe(
...     prompt="Palette knife painting of an autumn cityscape",  # What to generate
...     negative_prompt="Oversaturated, blurry, low quality",  # What NOT to generate
...     height=480,
...     width=640,  # Specify the image size
...     guidance_scale=8,  # How strongly to follow the prompt
...     num_inference_steps=35,  # How many steps to take
...     generator=generator,  # Fixed random seed
... )

>>> # View the resulting image
>>> pipe_output.images[0]

>>> # @markdown comparing guidance scales:
>>> cfg_scales = [1.1, 8, 12]  # @param
>>> prompt = "A collie with a pink hat"  # @param
>>> fig, axs = plt.subplots(1, len(cfg_scales), figsize=(16, 5))
>>> for i, ax in enumerate(axs):
...     im = pipe(
...         prompt,
...         height=480,
...         width=480,
...         guidance_scale=cfg_scales[i],
...         num_inference_steps=35,
...         generator=torch.Generator(device=device).manual_seed(42),
...     ).images[0]
...     ax.imshow(im)
...     ax.set_title(f"CFG Scale {cfg_scales[i]}")

>>> print(list(pipe.components.keys()))  # List components

>>> # Create some fake data (a random image, range (-1, 1))
>>> images = torch.rand(1, 3, 512, 512).to(device) * 2 - 1
>>> print("Input images shape:", images.shape)

>>> # Encode to latent space
>>> with torch.no_grad():
...     latents = 0.18215 * pipe.vae.encode(images).latent_dist.mean
>>> print("Encoded latents shape:", latents.shape)

>>> # Decode again
>>> with torch.no_grad():
...     decoded_images = pipe.vae.decode(latents / 0.18215).sample
>>> print("Decoded images shape:", decoded_images.shape)

>>> # Tokenizing and encoding an example prompt manually

>>> # Tokenize
>>> input_ids = pipe.tokenizer(["A painting of a flooble"])["input_ids"]
>>> print("Input ID -> decoded token")
>>> for input_id in input_ids[0]:
...     print(f"{input_id} -> {pipe.tokenizer.decode(input_id)}")

>>> # Feed through CLIP text encoder
>>> input_ids = torch.tensor(input_ids).to(device)
>>> with torch.no_grad():
...     text_embeddings = pipe.text_encoder(input_ids)["last_hidden_state"]
>>> print("Text embeddings shape:", text_embeddings.shape)

# Get the final text embeddings using the pipeline's encode_prompt function
text_embeddings = pipe._encode_prompt("A painting of a flooble", device, 1, False, "")
text_embeddings.shape

>>> # Dummy inputs
>>> timestep = pipe.scheduler.timesteps[0]
>>> latents = torch.randn(1, 4, 64, 64).to(device)
>>> text_embeddings = torch.randn(1, 77, 1024).to(device)

>>> # Model prediction
>>> with torch.no_grad():
...     unet_output = pipe.unet(latents, timestep, text_embeddings).sample
>>> print("UNet output shape:", unet_output.shape)  # Same shape as the input latents

>>> plt.plot(pipe.scheduler.alphas_cumprod, label=r"$\bar{\alpha}$")
>>> plt.xlabel("Timestep (high noise to low noise ->)")
>>> plt.title("Noise schedule")
>>> plt.legend()

>>> from diffusers import LMSDiscreteScheduler

>>> # Replace the scheduler
>>> pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)

>>> # Print the config
>>> print("Scheduler config:", pipe.scheduler)

>>> # Generate an image with this new scheduler
>>> pipe(
...     prompt="Palette knife painting of an winter cityscape",
...     height=480,
...     width=480,
...     generator=torch.Generator(device=device).manual_seed(42),
... ).images[0]

>>> guidance_scale = 8  # @param
>>> num_inference_steps = 30  # @param
>>> prompt = "Beautiful picture of a wave breaking"  # @param
>>> negative_prompt = "zoomed in, blurry, oversaturated, warped"  # @param

>>> # Encode the prompt
>>> text_embeddings = pipe._encode_prompt(prompt, device, 1, True, negative_prompt)

>>> # Create our random starting point
>>> latents = torch.randn((1, 4, 64, 64), device=device, generator=generator)
>>> latents *= pipe.scheduler.init_noise_sigma

>>> # Prepare the scheduler
>>> pipe.scheduler.set_timesteps(num_inference_steps, device=device)

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

...     # Expand the latents if we are doing classifier free guidance
...     latent_model_input = torch.cat([latents] * 2)

...     # Apply any scaling required by the scheduler
...     latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, t)

...     # Predict the noise residual with the UNet
...     with torch.no_grad():
...         noise_pred = pipe.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

...     # Perform guidance
...     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
...     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

...     # Compute the previous noisy sample x_t -> x_t-1
...     latents = pipe.scheduler.step(noise_pred, t, latents).prev_sample

>>> # Decode the resulting latents into an image
>>> with torch.no_grad():
...     image = pipe.decode_latents(latents.detach())

>>> # View
>>> pipe.numpy_to_pil(image)[0]

# Loading an Img2Img pipeline
model_id = "stabilityai/stable-diffusion-2-1-base"
img2img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id).to(device)

>>> # Apply Img2Img
>>> result_image = img2img_pipe(
...     prompt="An oil painting of a man on a bench",
...     image=init_image,  # The starting image
...     strength=0.6,  # 0 for no change, 1.0 for max strength
... ).images[0]

>>> # View the result
>>> fig, axs = plt.subplots(1, 2, figsize=(12, 5))
>>> axs[0].imshow(init_image)
>>> axs[0].set_title("Input Image")
>>> axs[1].imshow(result_image)
>>> axs[1].set_title("Result")

# Load the inpainting pipeline (requires a suitable inpainting model)
pipe = StableDiffusionInpaintPipeline.from_pretrained("runwayml/stable-diffusion-inpainting")
pipe = pipe.to(device)

>>> # Inpaint with a prompt for what we want the result to look like
>>> prompt = "A small robot, high resolution, sitting on a park bench"
>>> image = pipe(prompt=prompt, image=init_image, mask_image=mask_image).images[0]

>>> # View the result
>>> fig, axs = plt.subplots(1, 3, figsize=(16, 5))
>>> axs[0].imshow(init_image)
>>> axs[0].set_title("Input Image")
>>> axs[1].imshow(mask_image)
>>> axs[1].set_title("Mask")
>>> axs[2].imshow(image)
>>> axs[2].set_title("Result")

>>> !ls ~/.cache/huggingface/hub/ # List the contents of the cache directory

# Load the Depth2Img pipeline (requires a suitable model)
pipe = StableDiffusionDepth2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-depth")
pipe = pipe.to(device)

>>> # Inpaint with a prompt for what we want the result to look like
>>> prompt = "An oil painting of a man on a bench"
>>> image = pipe(prompt=prompt, image=init_image).images[0]

>>> # View the result
>>> fig, axs = plt.subplots(1, 2, figsize=(16, 5))
>>> axs[0].imshow(init_image)
>>> axs[0].set_title("Input Image")
>>> axs[1].imshow(image)
>>> axs[1].set_title("Result")