pip install diffusers transformers pillow torch datasets

import torch
from PIL import Image
import numpy as np
from datasets import load_dataset
from diffusers import DiffusionPipeline

# Load the CelebA-HQ dataset
ds = load_dataset("saitsharipov/CelebA-HQ", split="train")

# Get the first image from the dataset
image_data = ds[0]["image"]
image_pil = Image.fromarray(np.array(image_data))  # Convert to PIL image
image_name = "first_image"  # Use a descriptive name for saving files

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

# Load the diffusion pipeline
pipe = DiffusionPipeline.from_pretrained(
    "google/ddpm-ema-celebahq-256",
    custom_pipeline="ddim_noise_comparative_analysis",
)
pipe = pipe.to(device)

# Perform noise comparative analysis
for strength in np.linspace(0.1, 1, 25):
    denoised_image, latent_timestep = pipe(
        image_pil, strength=strength, return_dict=False
    )
    denoised_image = denoised_image[0]
    denoised_image.save(
        f"noise_comparative_analysis_{image_name}_{latent_timestep}.png"
    )