Blame for app.py

"""
Stable Diffusion Lab: a teaching Space for BA Design students.

Students can explore how the main knobs of Stable Diffusion 1.5 change an image:
prompt / negative prompt, CFG scale, number of denoising steps, sampler, checkpoint,
seed and resolution. While the image is being generated, a live preview of the
denoising process is streamed to the browser.

How the live preview works
--------------------------
We run our own (short) denoising loop instead of calling pipe(...). At every step the
U-Net predicts the noise. From that prediction we compute the model's current guess of
the final, clean latent ("x0 prediction") and decode it with TAESD, a tiny and very
fast VAE decoder. The last step is decoded with the full-quality VAE of the checkpoint.

Runs on a dedicated GPU (fp16) but also works on CPU for local testing (slow).
"""

import os
import random
import threading

import gradio as gr
import numpy as np
import torch
from diffusers import (
    AutoencoderKL,
    AutoencoderTiny,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    StableDiffusionPipeline,
    UniPCMultistepScheduler,
)
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from huggingface_hub import hf_hub_download
from PIL import Image
from transformers import CLIPImageProcessor

# ----------------------------------------------------------------------------
# Configuration (all overridable with Space variables / environment variables)
# ----------------------------------------------------------------------------
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = getattr(torch, os.getenv("DTYPE", "float16" if DEVICE == "cuda" else "float32"))

SD15_REPO = os.getenv("SD15_REPO", "sd-legacy/stable-diffusion-v1-5")
RV_REPO = os.getenv("RV_REPO", "SG161222/Realistic_Vision_V6.0_B1_noVAE")
RV_FILE = os.getenv("RV_FILE", "Realistic_Vision_V6.0_NV_B1_fp16.safetensors")
# Realistic Vision "noVAE" needs an external VAE (the author recommends sd-vae-ft-mse).
RV_VAE_REPO = os.getenv("RV_VAE_REPO", "stabilityai/sd-vae-ft-mse")
TAESD_REPO = os.getenv("TAESD_REPO", "madebyollin/taesd")
SAFETY_REPO = os.getenv("SAFETY_REPO", "CompVis/stable-diffusion-safety-checker")

USE_SAFETY_CHECKER = os.getenv("SAFETY_CHECKER", "1") == "1"
PRELOAD = os.getenv("PRELOAD", "1") == "1"  # load both checkpoints at startup
MAX_CONCURRENT = int(os.getenv("MAX_CONCURRENT", "1"))  # generations running at once

MODEL_NAMES = ["Stable Diffusion 1.5", "Realistic Vision V6.0"]
DEFAULT_MODEL = MODEL_NAMES[0]

SAMPLERS = {
    "Euler": EulerDiscreteScheduler,
    "Euler Ancestral": EulerAncestralDiscreteScheduler,
    "DDIM": DDIMScheduler,
    "UniPC": UniPCMultistepScheduler,
}

DEFAULTS = dict(steps=20, cfg=7.5, width=512, height=512, seed=42, sampler="Euler")

# ----------------------------------------------------------------------------
# Model loading
# ----------------------------------------------------------------------------
_pipes: dict = {}
_taesd = None
_safety = None
_load_lock = threading.Lock()


def _load_sd15() -> StableDiffusionPipeline:
    return StableDiffusionPipeline.from_pretrained(
        SD15_REPO,
        torch_dtype=DTYPE,
        variant="fp16",
        use_safetensors=True,
        safety_checker=None,  # we run the safety checker ourselves (see below)
        requires_safety_checker=False,
    )


def _load_realistic_vision() -> StableDiffusionPipeline:
    vae = AutoencoderKL.from_pretrained(RV_VAE_REPO, torch_dtype=DTYPE)
    ckpt = hf_hub_download(RV_REPO, RV_FILE)
    return StableDiffusionPipeline.from_single_file(
        ckpt,
        torch_dtype=DTYPE,
        vae=vae,
        config=SD15_REPO,  # SD 1.5 architecture, so the config files can come from there
        safety_checker=None,
        requires_safety_checker=False,
    )


_LOADERS = {
    "Stable Diffusion 1.5": _load_sd15,
    "Realistic Vision V6.0": _load_realistic_vision,
}


def get_pipe(name: str) -> StableDiffusionPipeline:
    with _load_lock:
        if name not in _pipes:
            print(f"Loading {name} ...")
            pipe = _LOADERS[name]().to(DEVICE)
            pipe.set_progress_bar_config(disable=True)
            _pipes[name] = pipe
        return _pipes[name]


def get_taesd() -> AutoencoderTiny:
    global _taesd
    with _load_lock:
        if _taesd is None:
            _taesd = AutoencoderTiny.from_pretrained(TAESD_REPO, torch_dtype=DTYPE).to(DEVICE)
        return _taesd


def get_safety():
    global _safety
    with _load_lock:
        if _safety is None:
            checker = StableDiffusionSafetyChecker.from_pretrained(SAFETY_REPO, torch_dtype=DTYPE).to(DEVICE)
            extractor = CLIPImageProcessor.from_pretrained(SAFETY_REPO)
            _safety = (checker, extractor)
        return _safety


# ----------------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------------
def to_pil(x: torch.Tensor) -> Image.Image:
    """(1, 3, H, W) tensor in [-1, 1] -> PIL image."""
    x = (x.float() / 2 + 0.5).clamp(0, 1)[0].permute(1, 2, 0).cpu().numpy()
    return Image.fromarray((x * 255).round().astype(np.uint8))


def predict_x0(scheduler, x_in: torch.Tensor, eps: torch.Tensor, t) -> torch.Tensor:
    """
    The model's current guess of the final clean latent, given the noisy latent and the
    predicted noise:  x0 = (x_t - sqrt(1 - abar) * eps) / sqrt(abar).

    `x_in` is the scheduler-scaled model input, i.e. the sample in "variance preserving"
    form. For the Euler samplers abar follows from sigma (abar = 1 / (sigma^2 + 1)),
    for DDIM and UniPC it is read from the training noise schedule.
    """
    if isinstance(scheduler, (EulerDiscreteScheduler, EulerAncestralDiscreteScheduler)):
        sigma = float(scheduler.sigmas[scheduler.step_index])
        abar = 1.0 / (sigma**2 + 1.0)
    else:
        abar = float(scheduler.alphas_cumprod[int(t)])
    return (x_in.float() - (1.0 - abar) ** 0.5 * eps.float()) / abar**0.5


def apply_safety_checker(image: Image.Image):
    """Returns (image, was_blocked)."""
    if not USE_SAFETY_CHECKER:
        return image, False
    checker, extractor = get_safety()
    clip_input = extractor(images=image, return_tensors="pt").pixel_values.to(DEVICE, DTYPE)
    arr = np.asarray(image).astype(np.float32)[None] / 255.0
    arr, flagged = checker(images=arr, clip_input=clip_input)
    if flagged and flagged[0]:
        return Image.fromarray((arr[0] * 255).astype(np.uint8)), True
    return image, False


# ----------------------------------------------------------------------------
# Generation (a generator: every yield updates the UI)
# ----------------------------------------------------------------------------
@torch.inference_mode()
def generate(model_name, prompt, negative_prompt, cfg, steps, sampler_name, seed, randomize_seed, width, height):
    steps, width, height = int(steps), int(width), int(height)
    seed = random.randint(0, 2**32 - 1) if randomize_seed or seed is None else int(seed)

    pipe = get_pipe(model_name)
    taesd = get_taesd()

    # A fresh scheduler per request: schedulers keep internal state, and several
    # students may be using the same pipeline.
    scheduler = SAMPLERS[sampler_name].from_config(pipe.scheduler.config, timestep_spacing="leading")
    scheduler.set_timesteps(steps, device=DEVICE)

    # Text conditioning. Batch order is [negative, positive] for classifier-free guidance.
    prompt_embeds, negative_embeds = pipe.encode_prompt(
        prompt or "", DEVICE, 1, True, negative_prompt=negative_prompt or ""
    )[:2]
    text_embeds = torch.cat([negative_embeds, prompt_embeds])

    # Seeded starting noise. Created on the CPU so a seed gives the same noise on any machine.
    generator = torch.Generator("cpu").manual_seed(seed)
    latents = torch.randn((1, pipe.unet.config.in_channels, height // 8, width // 8), generator=generator)
    latents = latents.to(DEVICE, DTYPE) * scheduler.init_noise_sigma

    step_kwargs = {"generator": generator} if isinstance(scheduler, EulerAncestralDiscreteScheduler) else {}

    for i, t in enumerate(scheduler.timesteps):
        x_in = scheduler.scale_model_input(torch.cat([latents, latents]), t)
        eps = pipe.unet(x_in, t, encoder_hidden_states=text_embeds).sample
        eps_uncond, eps_cond = eps.chunk(2)
        # Classifier-free guidance: 0 = ignore the prompt, 1 = plain prompt, >1 = push harder.
        eps = eps_uncond + cfg * (eps_cond - eps_uncond)

        if i < len(scheduler.timesteps) - 1:
            x0 = predict_x0(scheduler, x_in.chunk(2)[0], eps, t)
            preview = to_pil(taesd.decode(x0.to(DTYPE)).sample)
            latents = scheduler.step(eps, t, latents, **step_kwargs).prev_sample
            yield preview, seed, f"Step {i + 1} of {steps}: preview of the model's current guess"
        else:
            latents = scheduler.step(eps, t, latents, **step_kwargs).prev_sample

    # Final image with the full VAE of the checkpoint.
    vae = pipe.vae
    image = to_pil(vae.decode(latents / vae.config.scaling_factor).sample)
    image, blocked = apply_safety_checker(image)
    if blocked:
        gr.Warning("The safety filter blocked this image. Try a different prompt or seed.")
    yield image, seed, f"Done: {steps} steps, seed {seed}" + (" (blocked by safety filter)" if blocked else "")


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
EXAMPLES = [
    ["a cozy reading nook with a large window, soft morning light, plants, film photo", "blurry, low quality"],
    ["portrait of an old fisherman, weathered face, dramatic side light, 85mm photo", "cartoon, painting, deformed"],
    ["a futuristic tram station in the rain at night, neon reflections, wide angle", "text, watermark"],
    ["a ceramic teapot shaped like a mountain, studio product photo", ""],
]

with gr.Blocks(title="Stable Diffusion Lab") as demo:
    gr.Markdown(
        "# Stable Diffusion Lab\n"
        "Explore how text-to-image diffusion works. Change one setting at a time and keep the "
        "**seed** fixed to see exactly what that setting does. Watch the image emerge from noise "
        "in the live preview."
    )
    with gr.Row():
        with gr.Column(scale=5):
            model = gr.Dropdown(
                MODEL_NAMES,
                value=DEFAULT_MODEL,
                label="Checkpoint",
                info="Stable Diffusion 1.5 is the base model. Realistic Vision V6.0 is a finetune of it "
                "that was trained further towards photorealistic images.",
            )
            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="Describe what you want to see")
            negative = gr.Textbox(
                label="Negative prompt",
                lines=2,
                placeholder="What the image should avoid, e.g. blurry, low quality, extra fingers",
            )
            with gr.Row():
                sampler = gr.Dropdown(
                    list(SAMPLERS),
                    value=DEFAULTS["sampler"],
                    label="Sampler",
                    info="The algorithm that walks from noise to image. Euler Ancestral adds fresh noise "
                    "at every step, so it never fully settles.",
                )
                steps = gr.Slider(
                    4, 48, value=DEFAULTS["steps"], step=1, label="Denoising steps",
                    info="More steps = finer refinement, but slower.",
                )
            cfg = gr.Slider(
                0.0, 15.0, value=DEFAULTS["cfg"], step=0.1, label="CFG scale",
                info="How strongly the prompt is enforced. 0 ignores the prompt, 1 follows it plainly, "
                "high values get rigid and oversaturated.",
            )
            with gr.Row():
                width = gr.Slider(
                    256, 768, value=DEFAULTS["width"], step=64, label="Width (px)",
                    info="Both models were trained at 512 px. Much larger sizes tend to repeat objects.",
                )
                height = gr.Slider(256, 768, value=DEFAULTS["height"], step=64, label="Height (px)")
            with gr.Row():
                seed = gr.Number(
                    value=DEFAULTS["seed"], precision=0, minimum=0, maximum=2**32 - 1, label="Seed",
                    info="Same seed + same settings = same image.",
                )
                randomize = gr.Checkbox(value=False, label="Random seed", info="Pick a new seed on every run.")
            run = gr.Button("Generate", variant="primary")
        with gr.Column(scale=5):
            output = gr.Image(label="Result (live preview while generating)", type="pil", interactive=False)
            status = gr.Markdown()

    gr.Examples(EXAMPLES, inputs=[prompt, negative], label="Example prompts")

    inputs = [model, prompt, negative, cfg, steps, sampler, seed, randomize, width, height]
    outputs = [output, seed, status]
    for trigger in (run.click, prompt.submit):
        trigger(
            generate, inputs, outputs,
            show_progress="hidden", concurrency_limit=MAX_CONCURRENT, api_name="generate" if trigger == run.click else False,
        )

demo.queue(max_size=40)

if __name__ == "__main__":
    if PRELOAD:
        for _name in MODEL_NAMES:
            get_pipe(_name)
        get_taesd()
        if USE_SAFETY_CHECKER:
            get_safety()
    demo.launch()