Spaces:
Paused
Paused
Add Stable Diffusion equation playground
Browse files- README.md +41 -7
- app.py +607 -0
- requirements.txt +9 -0
README.md
CHANGED
|
@@ -1,13 +1,47 @@
|
|
| 1 |
---
|
| 2 |
-
title:
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: gradio
|
| 7 |
-
sdk_version:
|
| 8 |
-
python_version:
|
| 9 |
app_file: app.py
|
| 10 |
pinned: false
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 11 |
---
|
| 12 |
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
+
title: Stable Diffusion Equation Playground
|
| 3 |
+
emoji: 🧪
|
| 4 |
+
colorFrom: indigo
|
| 5 |
+
colorTo: green
|
| 6 |
sdk: gradio
|
| 7 |
+
sdk_version: 5.22.0
|
| 8 |
+
python_version: 3.10
|
| 9 |
app_file: app.py
|
| 10 |
pinned: false
|
| 11 |
+
license: mit
|
| 12 |
+
models:
|
| 13 |
+
- stable-diffusion-v1-5/stable-diffusion-v1-5
|
| 14 |
+
preload_from_hub:
|
| 15 |
+
- stable-diffusion-v1-5/stable-diffusion-v1-5
|
| 16 |
---
|
| 17 |
|
| 18 |
+
# Stable Diffusion Equation Playground
|
| 19 |
+
|
| 20 |
+
A ZeroGPU-ready Gradio app for learning how Stable Diffusion works inside Diffusers.
|
| 21 |
+
|
| 22 |
+
Instead of only calling `pipe(prompt)`, the app exposes a small custom denoising loop:
|
| 23 |
+
|
| 24 |
+
- prompt embeddings can be used directly, averaged, or combined with vector arithmetic
|
| 25 |
+
- initial latent noise can come from one seed or a blend of two seeds
|
| 26 |
+
- classifier-free guidance can use standard CFG or a student-edited equation
|
| 27 |
+
- intermediate latent snapshots show how the image emerges across denoising steps
|
| 28 |
+
|
| 29 |
+
The core idea to teach: Stable Diffusion starts from noise in VAE latent space. The prompt does not literally become half of the pixels. The prompt changes the UNet's predicted noise at every denoising step, most commonly through classifier-free guidance:
|
| 30 |
+
|
| 31 |
+
```python
|
| 32 |
+
guided = negative_prediction + guidance_scale * (prompt_prediction - negative_prediction)
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
This app focuses on `StableDiffusionPipeline` checkpoints such as Stable Diffusion 1.5 because the embedding pathway is straightforward for students. SDXL uses two text encoders plus pooled prompt embeddings, so the same idea carries over, but the code is more complex.
|
| 36 |
+
|
| 37 |
+
## Running
|
| 38 |
+
|
| 39 |
+
On Hugging Face Spaces, select ZeroGPU hardware in the Space settings. The image generation function is decorated with `@spaces.GPU`, so it requests a GPU only while a generation is running.
|
| 40 |
+
|
| 41 |
+
For local development, install the requirements, use a CUDA or MPS machine, then run:
|
| 42 |
+
|
| 43 |
+
```bash
|
| 44 |
+
python app.py
|
| 45 |
+
```
|
| 46 |
+
|
| 47 |
+
The default model is `stable-diffusion-v1-5/stable-diffusion-v1-5`. The first generation downloads the checkpoint from Hugging Face and may take a few minutes.
|
app.py
ADDED
|
@@ -0,0 +1,607 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import gc
|
| 2 |
+
from functools import lru_cache
|
| 3 |
+
|
| 4 |
+
import gradio as gr
|
| 5 |
+
import numpy as np
|
| 6 |
+
import torch
|
| 7 |
+
from PIL import Image, ImageDraw
|
| 8 |
+
|
| 9 |
+
from diffusers import DDIMScheduler, DPMSolverMultistepScheduler, StableDiffusionPipeline
|
| 10 |
+
|
| 11 |
+
try:
|
| 12 |
+
import spaces
|
| 13 |
+
except ImportError:
|
| 14 |
+
class _SpacesFallback:
|
| 15 |
+
@staticmethod
|
| 16 |
+
def GPU(*decorator_args, **decorator_kwargs):
|
| 17 |
+
if decorator_args and callable(decorator_args[0]) and not decorator_kwargs:
|
| 18 |
+
return decorator_args[0]
|
| 19 |
+
|
| 20 |
+
def decorator(func):
|
| 21 |
+
return func
|
| 22 |
+
|
| 23 |
+
return decorator
|
| 24 |
+
|
| 25 |
+
spaces = _SpacesFallback()
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
APP_TITLE = "Stable Diffusion Equation Playground"
|
| 29 |
+
DEFAULT_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5"
|
| 30 |
+
DEFAULT_PROMPT_A = "a small glass greenhouse filled with glowing ferns, watercolor"
|
| 31 |
+
DEFAULT_PROMPT_B = "a futuristic train station at sunrise, cinematic lighting"
|
| 32 |
+
DEFAULT_PROMPT_C = "low quality, blurry, distorted"
|
| 33 |
+
MAX_SEED = 2_147_483_647
|
| 34 |
+
|
| 35 |
+
PROMPT_MATH_CODE = """# Diffusers usually does this inside pipe(prompt).
|
| 36 |
+
# Here we expose the text embeddings so we can edit them directly.
|
| 37 |
+
prompt_a, negative = encode_prompt(prompt_a, negative_prompt)
|
| 38 |
+
prompt_b, _ = encode_prompt(prompt_b, negative_prompt)
|
| 39 |
+
prompt_c, _ = encode_prompt(prompt_c, negative_prompt)
|
| 40 |
+
|
| 41 |
+
if mode == "average":
|
| 42 |
+
prompt_embeds = (1 - mix) * prompt_a + mix * prompt_b
|
| 43 |
+
elif mode == "analogy":
|
| 44 |
+
prompt_embeds = prompt_a + strength * (prompt_b - prompt_c)
|
| 45 |
+
else:
|
| 46 |
+
prompt_embeds = prompt_a
|
| 47 |
+
"""
|
| 48 |
+
|
| 49 |
+
LATENT_MATH_CODE = """# Stable Diffusion does not start from pixels.
|
| 50 |
+
# It starts from noisy latents in the VAE's compressed image space.
|
| 51 |
+
noise_a = torch.randn(latent_shape, generator=seed_a)
|
| 52 |
+
noise_b = torch.randn(latent_shape, generator=seed_b)
|
| 53 |
+
latents = (1 - noise_mix) * noise_a + noise_mix * noise_b
|
| 54 |
+
|
| 55 |
+
if renormalize_noise:
|
| 56 |
+
latents = (latents - latents.mean()) / latents.std()
|
| 57 |
+
|
| 58 |
+
latents = latents * scheduler.init_noise_sigma
|
| 59 |
+
"""
|
| 60 |
+
|
| 61 |
+
GUIDANCE_MATH_CODE = """# Classifier-free guidance combines two UNet predictions:
|
| 62 |
+
# one conditioned on the negative/unconditional prompt, one on the prompt.
|
| 63 |
+
noise_negative, noise_prompt = noise_pred.chunk(2)
|
| 64 |
+
delta = noise_prompt - noise_negative
|
| 65 |
+
|
| 66 |
+
# Standard CFG is:
|
| 67 |
+
# guided = noise_negative + guidance_scale * delta
|
| 68 |
+
guided = negative_coeff * noise_negative
|
| 69 |
+
guided = guided + prompt_coeff * noise_prompt
|
| 70 |
+
guided = guided + delta_coeff * delta
|
| 71 |
+
"""
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def current_device():
|
| 75 |
+
if torch.cuda.is_available():
|
| 76 |
+
return torch.device("cuda")
|
| 77 |
+
if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
|
| 78 |
+
return torch.device("mps")
|
| 79 |
+
return torch.device("cpu")
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def model_dtype(device):
|
| 83 |
+
if device.type == "cuda":
|
| 84 |
+
return torch.float16
|
| 85 |
+
return torch.float32
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
def device_label(device):
|
| 89 |
+
if device.type == "cuda":
|
| 90 |
+
name = torch.cuda.get_device_name(0)
|
| 91 |
+
return f"CUDA GPU: {name}"
|
| 92 |
+
if device.type == "mps":
|
| 93 |
+
return "Apple MPS GPU"
|
| 94 |
+
return "CPU only. This app will load, but image generation will be very slow."
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def round_to_multiple_of_8(value):
|
| 98 |
+
value = int(value)
|
| 99 |
+
return max(256, min(768, 8 * round(value / 8)))
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
def seed_generator(seed, device):
|
| 103 |
+
seed = int(seed) % MAX_SEED
|
| 104 |
+
if device.type == "cuda":
|
| 105 |
+
return torch.Generator(device=device).manual_seed(seed)
|
| 106 |
+
return torch.Generator(device="cpu").manual_seed(seed)
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def randn_tensor(shape, seed, device, dtype):
|
| 110 |
+
generator = seed_generator(seed, device)
|
| 111 |
+
if device.type == "cuda":
|
| 112 |
+
return torch.randn(shape, generator=generator, device=device, dtype=dtype)
|
| 113 |
+
return torch.randn(shape, generator=generator, device="cpu", dtype=dtype).to(device)
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def blank_image(message="Run generation to make an image."):
|
| 117 |
+
image = Image.new("RGB", (512, 512), (25, 29, 36))
|
| 118 |
+
draw = ImageDraw.Draw(image)
|
| 119 |
+
draw.text((32, 236), message, fill=(230, 235, 240))
|
| 120 |
+
return image
|
| 121 |
+
|
| 122 |
+
|
| 123 |
+
@lru_cache(maxsize=2)
|
| 124 |
+
def load_pipe(model_id, scheduler_name, device_type):
|
| 125 |
+
device = torch.device(device_type)
|
| 126 |
+
dtype = model_dtype(device)
|
| 127 |
+
pipe = StableDiffusionPipeline.from_pretrained(
|
| 128 |
+
model_id,
|
| 129 |
+
torch_dtype=dtype,
|
| 130 |
+
use_safetensors=True,
|
| 131 |
+
)
|
| 132 |
+
|
| 133 |
+
if scheduler_name == "DDIM":
|
| 134 |
+
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
|
| 135 |
+
else:
|
| 136 |
+
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
| 137 |
+
|
| 138 |
+
pipe = pipe.to(device)
|
| 139 |
+
pipe.set_progress_bar_config(disable=True)
|
| 140 |
+
pipe.enable_vae_slicing()
|
| 141 |
+
|
| 142 |
+
if device.type == "cuda":
|
| 143 |
+
try:
|
| 144 |
+
pipe.enable_xformers_memory_efficient_attention()
|
| 145 |
+
except Exception:
|
| 146 |
+
pass
|
| 147 |
+
|
| 148 |
+
return pipe
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def encode_prompt(pipe, prompt, negative_prompt, device):
|
| 152 |
+
if hasattr(pipe, "encode_prompt"):
|
| 153 |
+
prompt_embeds, negative_prompt_embeds = pipe.encode_prompt(
|
| 154 |
+
prompt=prompt,
|
| 155 |
+
device=device,
|
| 156 |
+
num_images_per_prompt=1,
|
| 157 |
+
do_classifier_free_guidance=True,
|
| 158 |
+
negative_prompt=negative_prompt,
|
| 159 |
+
)
|
| 160 |
+
return prompt_embeds, negative_prompt_embeds
|
| 161 |
+
|
| 162 |
+
combined = pipe._encode_prompt(
|
| 163 |
+
prompt=prompt,
|
| 164 |
+
device=device,
|
| 165 |
+
num_images_per_prompt=1,
|
| 166 |
+
do_classifier_free_guidance=True,
|
| 167 |
+
negative_prompt=negative_prompt,
|
| 168 |
+
)
|
| 169 |
+
negative_prompt_embeds, prompt_embeds = combined.chunk(2)
|
| 170 |
+
return prompt_embeds, negative_prompt_embeds
|
| 171 |
+
|
| 172 |
+
|
| 173 |
+
def cosine_similarity(a, b):
|
| 174 |
+
a = a.detach().float().flatten()
|
| 175 |
+
b = b.detach().float().flatten()
|
| 176 |
+
return float(torch.nn.functional.cosine_similarity(a, b, dim=0).cpu())
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def mix_prompt_embeddings(
|
| 180 |
+
pipe,
|
| 181 |
+
device,
|
| 182 |
+
prompt_a,
|
| 183 |
+
prompt_b,
|
| 184 |
+
prompt_c,
|
| 185 |
+
negative_prompt,
|
| 186 |
+
embedding_mode,
|
| 187 |
+
prompt_mix,
|
| 188 |
+
analogy_strength,
|
| 189 |
+
renormalize_prompt,
|
| 190 |
+
):
|
| 191 |
+
emb_a, negative_embeds = encode_prompt(pipe, prompt_a, negative_prompt, device)
|
| 192 |
+
emb_b, _ = encode_prompt(pipe, prompt_b, negative_prompt, device)
|
| 193 |
+
emb_c, _ = encode_prompt(pipe, prompt_c or "", negative_prompt, device)
|
| 194 |
+
|
| 195 |
+
if embedding_mode == "Average prompt A and B":
|
| 196 |
+
mixed = (1.0 - prompt_mix) * emb_a + prompt_mix * emb_b
|
| 197 |
+
formula = f"prompt = {(1.0 - prompt_mix):.2f} * A + {prompt_mix:.2f} * B"
|
| 198 |
+
elif embedding_mode == "Vector arithmetic: A + s * (B - C)":
|
| 199 |
+
mixed = emb_a + analogy_strength * (emb_b - emb_c)
|
| 200 |
+
formula = f"prompt = A + {analogy_strength:.2f} * (B - C)"
|
| 201 |
+
else:
|
| 202 |
+
mixed = emb_a
|
| 203 |
+
formula = "prompt = A"
|
| 204 |
+
|
| 205 |
+
original_norm = emb_a.detach().float().norm()
|
| 206 |
+
mixed_norm = mixed.detach().float().norm()
|
| 207 |
+
if renormalize_prompt and float(mixed_norm.cpu()) > 0:
|
| 208 |
+
mixed = mixed * (original_norm / mixed_norm)
|
| 209 |
+
formula += "; then rescale to A's embedding norm"
|
| 210 |
+
|
| 211 |
+
metrics = [
|
| 212 |
+
["cosine(A, B)", round(cosine_similarity(emb_a, emb_b), 4)],
|
| 213 |
+
["cosine(A, mixed)", round(cosine_similarity(emb_a, mixed), 4)],
|
| 214 |
+
["cosine(B, mixed)", round(cosine_similarity(emb_b, mixed), 4)],
|
| 215 |
+
["norm(A)", round(float(original_norm.cpu()), 3)],
|
| 216 |
+
["norm(mixed)", round(float(mixed.detach().float().norm().cpu()), 3)],
|
| 217 |
+
]
|
| 218 |
+
return mixed, negative_embeds, formula, metrics
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def prepare_latents(pipe, device, height, width, seed_a, seed_b, noise_mix, renormalize_noise):
|
| 222 |
+
channels = int(pipe.unet.config.in_channels)
|
| 223 |
+
latent_shape = (1, channels, height // pipe.vae_scale_factor, width // pipe.vae_scale_factor)
|
| 224 |
+
dtype = model_dtype(device)
|
| 225 |
+
noise_a = randn_tensor(latent_shape, seed_a, device, dtype)
|
| 226 |
+
noise_b = randn_tensor(latent_shape, seed_b, device, dtype)
|
| 227 |
+
latents = (1.0 - noise_mix) * noise_a + noise_mix * noise_b
|
| 228 |
+
|
| 229 |
+
before_std = float(latents.detach().float().std().cpu())
|
| 230 |
+
if renormalize_noise:
|
| 231 |
+
latents = (latents - latents.mean()) / (latents.std() + 1e-6)
|
| 232 |
+
after_std = float(latents.detach().float().std().cpu())
|
| 233 |
+
|
| 234 |
+
latents = latents * pipe.scheduler.init_noise_sigma
|
| 235 |
+
formula = f"noise = {(1.0 - noise_mix):.2f} * seed A + {noise_mix:.2f} * seed B"
|
| 236 |
+
if renormalize_noise:
|
| 237 |
+
formula += "; then renormalize to unit standard deviation"
|
| 238 |
+
|
| 239 |
+
metrics = [
|
| 240 |
+
["latent shape", str(tuple(latent_shape))],
|
| 241 |
+
["std before scheduler scale", round(before_std, 4)],
|
| 242 |
+
["std after optional renorm", round(after_std, 4)],
|
| 243 |
+
["scheduler init sigma", round(float(pipe.scheduler.init_noise_sigma), 4)],
|
| 244 |
+
]
|
| 245 |
+
return latents, formula, metrics
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
def combine_noise_predictions(
|
| 249 |
+
noise_negative,
|
| 250 |
+
noise_prompt,
|
| 251 |
+
equation_mode,
|
| 252 |
+
guidance_scale,
|
| 253 |
+
prediction_mix,
|
| 254 |
+
negative_coeff,
|
| 255 |
+
prompt_coeff,
|
| 256 |
+
delta_coeff,
|
| 257 |
+
):
|
| 258 |
+
delta = noise_prompt - noise_negative
|
| 259 |
+
|
| 260 |
+
if equation_mode == "Standard CFG: negative + scale * (prompt - negative)":
|
| 261 |
+
guided = noise_negative + guidance_scale * delta
|
| 262 |
+
formula = f"guided = negative + {guidance_scale:.2f} * (prompt - negative)"
|
| 263 |
+
elif equation_mode == "Direct blend: (1 - mix) * negative + mix * prompt":
|
| 264 |
+
guided = (1.0 - prediction_mix) * noise_negative + prediction_mix * noise_prompt
|
| 265 |
+
formula = f"guided = {(1.0 - prediction_mix):.2f} * negative + {prediction_mix:.2f} * prompt"
|
| 266 |
+
elif equation_mode == "Prompt prediction only":
|
| 267 |
+
guided = noise_prompt
|
| 268 |
+
formula = "guided = prompt"
|
| 269 |
+
elif equation_mode == "Negative prediction only":
|
| 270 |
+
guided = noise_negative
|
| 271 |
+
formula = "guided = negative"
|
| 272 |
+
else:
|
| 273 |
+
guided = negative_coeff * noise_negative + prompt_coeff * noise_prompt + delta_coeff * delta
|
| 274 |
+
formula = (
|
| 275 |
+
f"guided = {negative_coeff:.2f} * negative + {prompt_coeff:.2f} * prompt "
|
| 276 |
+
f"+ {delta_coeff:.2f} * (prompt - negative)"
|
| 277 |
+
)
|
| 278 |
+
|
| 279 |
+
return guided, formula
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
def decode_latents(pipe, latents):
|
| 283 |
+
latents = latents / pipe.vae.config.scaling_factor
|
| 284 |
+
image = pipe.vae.decode(latents, return_dict=False)[0]
|
| 285 |
+
image = (image / 2 + 0.5).clamp(0, 1)
|
| 286 |
+
image = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
|
| 287 |
+
image = (image * 255).round().astype("uint8")
|
| 288 |
+
return [Image.fromarray(frame) for frame in image]
|
| 289 |
+
|
| 290 |
+
|
| 291 |
+
def checkpoint_indices(num_steps):
|
| 292 |
+
last = max(0, int(num_steps) - 1)
|
| 293 |
+
return sorted({0, last // 3, (2 * last) // 3, last})
|
| 294 |
+
|
| 295 |
+
|
| 296 |
+
def gpu_duration(*args):
|
| 297 |
+
try:
|
| 298 |
+
steps = int(args[-3])
|
| 299 |
+
width = int(args[-2])
|
| 300 |
+
height = int(args[-1])
|
| 301 |
+
except Exception:
|
| 302 |
+
return 90
|
| 303 |
+
|
| 304 |
+
pixel_factor = max(1.0, (width * height) / (512 * 512))
|
| 305 |
+
return min(180, max(60, int(35 + steps * 2.5 * pixel_factor)))
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
@spaces.GPU(duration=gpu_duration)
|
| 309 |
+
@torch.inference_mode()
|
| 310 |
+
def generate(
|
| 311 |
+
model_id,
|
| 312 |
+
scheduler_name,
|
| 313 |
+
prompt_a,
|
| 314 |
+
prompt_b,
|
| 315 |
+
prompt_c,
|
| 316 |
+
negative_prompt,
|
| 317 |
+
embedding_mode,
|
| 318 |
+
prompt_mix,
|
| 319 |
+
analogy_strength,
|
| 320 |
+
renormalize_prompt,
|
| 321 |
+
seed_a,
|
| 322 |
+
seed_b,
|
| 323 |
+
noise_mix,
|
| 324 |
+
renormalize_noise,
|
| 325 |
+
equation_mode,
|
| 326 |
+
guidance_scale,
|
| 327 |
+
prediction_mix,
|
| 328 |
+
negative_coeff,
|
| 329 |
+
prompt_coeff,
|
| 330 |
+
delta_coeff,
|
| 331 |
+
num_steps,
|
| 332 |
+
width,
|
| 333 |
+
height,
|
| 334 |
+
):
|
| 335 |
+
device = current_device()
|
| 336 |
+
if device.type == "cpu":
|
| 337 |
+
return (
|
| 338 |
+
blank_image("GPU recommended."),
|
| 339 |
+
[],
|
| 340 |
+
"No GPU was detected. The app is designed for CUDA or MPS. It can run on CPU, but it may take a very long time.",
|
| 341 |
+
[],
|
| 342 |
+
[],
|
| 343 |
+
)
|
| 344 |
+
|
| 345 |
+
width = round_to_multiple_of_8(width)
|
| 346 |
+
height = round_to_multiple_of_8(height)
|
| 347 |
+
num_steps = int(num_steps)
|
| 348 |
+
pipe = load_pipe(model_id.strip() or DEFAULT_MODEL, scheduler_name, device.type)
|
| 349 |
+
|
| 350 |
+
prompt_embeds, negative_prompt_embeds, prompt_formula, prompt_metrics = mix_prompt_embeddings(
|
| 351 |
+
pipe,
|
| 352 |
+
device,
|
| 353 |
+
prompt_a or "",
|
| 354 |
+
prompt_b or "",
|
| 355 |
+
prompt_c or "",
|
| 356 |
+
negative_prompt or "",
|
| 357 |
+
embedding_mode,
|
| 358 |
+
float(prompt_mix),
|
| 359 |
+
float(analogy_strength),
|
| 360 |
+
bool(renormalize_prompt),
|
| 361 |
+
)
|
| 362 |
+
|
| 363 |
+
pipe.scheduler.set_timesteps(num_steps, device=device)
|
| 364 |
+
latents, latent_formula, latent_metrics = prepare_latents(
|
| 365 |
+
pipe,
|
| 366 |
+
device,
|
| 367 |
+
height,
|
| 368 |
+
width,
|
| 369 |
+
int(seed_a),
|
| 370 |
+
int(seed_b),
|
| 371 |
+
float(noise_mix),
|
| 372 |
+
bool(renormalize_noise),
|
| 373 |
+
)
|
| 374 |
+
|
| 375 |
+
text_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
|
| 376 |
+
snapshots = []
|
| 377 |
+
trace_rows = []
|
| 378 |
+
save_at = checkpoint_indices(num_steps)
|
| 379 |
+
last_formula = ""
|
| 380 |
+
|
| 381 |
+
for step_index, timestep in enumerate(pipe.scheduler.timesteps):
|
| 382 |
+
latent_model_input = torch.cat([latents] * 2)
|
| 383 |
+
latent_model_input = pipe.scheduler.scale_model_input(latent_model_input, timestep)
|
| 384 |
+
|
| 385 |
+
noise_pred = pipe.unet(
|
| 386 |
+
latent_model_input,
|
| 387 |
+
timestep,
|
| 388 |
+
encoder_hidden_states=text_embeds,
|
| 389 |
+
return_dict=False,
|
| 390 |
+
)[0]
|
| 391 |
+
noise_negative, noise_prompt = noise_pred.chunk(2)
|
| 392 |
+
guided, last_formula = combine_noise_predictions(
|
| 393 |
+
noise_negative,
|
| 394 |
+
noise_prompt,
|
| 395 |
+
equation_mode,
|
| 396 |
+
float(guidance_scale),
|
| 397 |
+
float(prediction_mix),
|
| 398 |
+
float(negative_coeff),
|
| 399 |
+
float(prompt_coeff),
|
| 400 |
+
float(delta_coeff),
|
| 401 |
+
)
|
| 402 |
+
|
| 403 |
+
latents = pipe.scheduler.step(guided, timestep, latents, return_dict=False)[0]
|
| 404 |
+
|
| 405 |
+
delta_norm = float((noise_prompt - noise_negative).detach().float().norm().cpu())
|
| 406 |
+
guided_norm = float(guided.detach().float().norm().cpu())
|
| 407 |
+
trace_rows.append(
|
| 408 |
+
[
|
| 409 |
+
step_index + 1,
|
| 410 |
+
int(timestep.detach().cpu()) if torch.is_tensor(timestep) else int(timestep),
|
| 411 |
+
round(delta_norm, 3),
|
| 412 |
+
round(guided_norm, 3),
|
| 413 |
+
]
|
| 414 |
+
)
|
| 415 |
+
|
| 416 |
+
if step_index in save_at:
|
| 417 |
+
snapshot = decode_latents(pipe, latents)[0]
|
| 418 |
+
snapshots.append((snapshot, f"step {step_index + 1} of {num_steps}"))
|
| 419 |
+
|
| 420 |
+
final_image = decode_latents(pipe, latents)[0]
|
| 421 |
+
metrics = prompt_metrics + latent_metrics + [
|
| 422 |
+
["device", device_label(device)],
|
| 423 |
+
["prompt formula", prompt_formula],
|
| 424 |
+
["noise formula", latent_formula],
|
| 425 |
+
["guidance formula", last_formula],
|
| 426 |
+
]
|
| 427 |
+
summary = (
|
| 428 |
+
f"Model: {model_id.strip() or DEFAULT_MODEL}\n"
|
| 429 |
+
f"Scheduler: {scheduler_name}; steps: {num_steps}; size: {width}x{height}\n\n"
|
| 430 |
+
f"Prompt embedding math:\n{prompt_formula}\n\n"
|
| 431 |
+
f"Initial latent noise math:\n{latent_formula}\n\n"
|
| 432 |
+
f"Guidance equation used at each denoising step:\n{last_formula}\n\n"
|
| 433 |
+
"Stable Diffusion starts from noise in latent space. The prompt does not get averaged with pixels; "
|
| 434 |
+
"it changes the UNet's predicted noise at every step."
|
| 435 |
+
)
|
| 436 |
+
|
| 437 |
+
if device.type == "cuda":
|
| 438 |
+
torch.cuda.empty_cache()
|
| 439 |
+
gc.collect()
|
| 440 |
+
return final_image, snapshots, summary, metrics, trace_rows
|
| 441 |
+
|
| 442 |
+
|
| 443 |
+
def randomize_seeds():
|
| 444 |
+
rng = np.random.default_rng()
|
| 445 |
+
return int(rng.integers(0, MAX_SEED)), int(rng.integers(0, MAX_SEED))
|
| 446 |
+
|
| 447 |
+
|
| 448 |
+
def build_app():
|
| 449 |
+
theme = gr.themes.Soft(
|
| 450 |
+
primary_hue="indigo",
|
| 451 |
+
secondary_hue="emerald",
|
| 452 |
+
neutral_hue="slate",
|
| 453 |
+
radius_size="sm",
|
| 454 |
+
)
|
| 455 |
+
|
| 456 |
+
css = """
|
| 457 |
+
.snapshot-gallery img { object-fit: contain !important; }
|
| 458 |
+
.code-panel textarea, .code-panel pre { font-size: 13px !important; }
|
| 459 |
+
"""
|
| 460 |
+
|
| 461 |
+
metric_headers = ["quantity", "value"]
|
| 462 |
+
trace_headers = ["step", "timestep", "prompt minus negative norm", "guided prediction norm"]
|
| 463 |
+
|
| 464 |
+
with gr.Blocks(title=APP_TITLE, theme=theme, css=css) as demo:
|
| 465 |
+
gr.Markdown(
|
| 466 |
+
f"# {APP_TITLE}\n"
|
| 467 |
+
"Change the starting noise, average prompt embeddings, and edit the classifier-free guidance equation. "
|
| 468 |
+
"This is a GPU-oriented teaching app built on Diffusers."
|
| 469 |
+
)
|
| 470 |
+
|
| 471 |
+
with gr.Row(equal_height=False):
|
| 472 |
+
with gr.Column(scale=1, min_width=320):
|
| 473 |
+
model_id = gr.Textbox(value=DEFAULT_MODEL, label="Diffusers model id")
|
| 474 |
+
scheduler_name = gr.Radio(["DPM++ 2M", "DDIM"], value="DPM++ 2M", label="Scheduler")
|
| 475 |
+
|
| 476 |
+
with gr.Accordion("Prompt embeddings", open=True):
|
| 477 |
+
prompt_a = gr.Textbox(value=DEFAULT_PROMPT_A, label="Prompt A", lines=3)
|
| 478 |
+
prompt_b = gr.Textbox(value=DEFAULT_PROMPT_B, label="Prompt B", lines=3)
|
| 479 |
+
prompt_c = gr.Textbox(value=DEFAULT_PROMPT_C, label="Prompt C for vector arithmetic", lines=2)
|
| 480 |
+
negative_prompt = gr.Textbox(
|
| 481 |
+
value="blurry, low quality, distorted",
|
| 482 |
+
label="Negative prompt",
|
| 483 |
+
lines=2,
|
| 484 |
+
)
|
| 485 |
+
embedding_mode = gr.Dropdown(
|
| 486 |
+
[
|
| 487 |
+
"Prompt A only",
|
| 488 |
+
"Average prompt A and B",
|
| 489 |
+
"Vector arithmetic: A + s * (B - C)",
|
| 490 |
+
],
|
| 491 |
+
value="Average prompt A and B",
|
| 492 |
+
label="Embedding equation",
|
| 493 |
+
)
|
| 494 |
+
prompt_mix = gr.Slider(0, 1, value=0.5, step=0.05, label="Prompt B weight")
|
| 495 |
+
analogy_strength = gr.Slider(-2, 2, value=0.8, step=0.1, label="Vector arithmetic strength")
|
| 496 |
+
renormalize_prompt = gr.Checkbox(value=True, label="Keep mixed prompt embedding norm near prompt A")
|
| 497 |
+
|
| 498 |
+
with gr.Accordion("Starting latent noise", open=True):
|
| 499 |
+
with gr.Row():
|
| 500 |
+
seed_a = gr.Number(value=11, precision=0, label="Seed A")
|
| 501 |
+
seed_b = gr.Number(value=2222, precision=0, label="Seed B")
|
| 502 |
+
random_seeds = gr.Button("Randomize seeds")
|
| 503 |
+
noise_mix = gr.Slider(0, 1, value=0.0, step=0.05, label="Seed B noise weight")
|
| 504 |
+
renormalize_noise = gr.Checkbox(value=True, label="Renormalize mixed noise")
|
| 505 |
+
|
| 506 |
+
with gr.Accordion("Guidance equation", open=True):
|
| 507 |
+
equation_mode = gr.Dropdown(
|
| 508 |
+
[
|
| 509 |
+
"Standard CFG: negative + scale * (prompt - negative)",
|
| 510 |
+
"Direct blend: (1 - mix) * negative + mix * prompt",
|
| 511 |
+
"Prompt prediction only",
|
| 512 |
+
"Negative prediction only",
|
| 513 |
+
"Custom coefficients",
|
| 514 |
+
],
|
| 515 |
+
value="Standard CFG: negative + scale * (prompt - negative)",
|
| 516 |
+
label="Noise prediction equation",
|
| 517 |
+
)
|
| 518 |
+
guidance_scale = gr.Slider(0, 15, value=7.5, step=0.25, label="CFG scale")
|
| 519 |
+
prediction_mix = gr.Slider(0, 2, value=0.5, step=0.05, label="Direct blend prompt weight")
|
| 520 |
+
with gr.Row():
|
| 521 |
+
negative_coeff = gr.Slider(-3, 3, value=1, step=0.1, label="negative coeff")
|
| 522 |
+
prompt_coeff = gr.Slider(-3, 3, value=0, step=0.1, label="prompt coeff")
|
| 523 |
+
delta_coeff = gr.Slider(-3, 15, value=7.5, step=0.25, label="delta coeff")
|
| 524 |
+
|
| 525 |
+
with gr.Accordion("Render settings", open=True):
|
| 526 |
+
with gr.Row():
|
| 527 |
+
width = gr.Slider(256, 768, value=512, step=8, label="Width")
|
| 528 |
+
height = gr.Slider(256, 768, value=512, step=8, label="Height")
|
| 529 |
+
num_steps = gr.Slider(4, 50, value=20, step=1, label="Denoising steps")
|
| 530 |
+
generate_button = gr.Button("Generate", variant="primary")
|
| 531 |
+
|
| 532 |
+
with gr.Column(scale=1, min_width=420):
|
| 533 |
+
output_image = gr.Image(
|
| 534 |
+
value=blank_image(),
|
| 535 |
+
label="Generated image",
|
| 536 |
+
type="pil",
|
| 537 |
+
interactive=False,
|
| 538 |
+
)
|
| 539 |
+
summary = gr.Textbox(label="What happened", lines=12, interactive=False)
|
| 540 |
+
snapshots = gr.Gallery(
|
| 541 |
+
label="Decoded latent snapshots",
|
| 542 |
+
columns=2,
|
| 543 |
+
height=420,
|
| 544 |
+
object_fit="contain",
|
| 545 |
+
elem_classes=["snapshot-gallery"],
|
| 546 |
+
)
|
| 547 |
+
metrics = gr.Dataframe(
|
| 548 |
+
headers=metric_headers,
|
| 549 |
+
datatype=["str", "str"],
|
| 550 |
+
label="Embedding and latent measurements",
|
| 551 |
+
interactive=False,
|
| 552 |
+
)
|
| 553 |
+
trace = gr.Dataframe(
|
| 554 |
+
headers=trace_headers,
|
| 555 |
+
datatype=["number", "number", "number", "number"],
|
| 556 |
+
label="Denoising trace",
|
| 557 |
+
interactive=False,
|
| 558 |
+
)
|
| 559 |
+
|
| 560 |
+
with gr.Tab("Code cells"):
|
| 561 |
+
with gr.Row(equal_height=False):
|
| 562 |
+
gr.Code(PROMPT_MATH_CODE, language="python", label="Prompt embedding math", interactive=False, elem_classes=["code-panel"])
|
| 563 |
+
gr.Code(LATENT_MATH_CODE, language="python", label="Latent noise math", interactive=False, elem_classes=["code-panel"])
|
| 564 |
+
gr.Code(GUIDANCE_MATH_CODE, language="python", label="Guidance equation", interactive=False, elem_classes=["code-panel"])
|
| 565 |
+
|
| 566 |
+
random_seeds.click(
|
| 567 |
+
randomize_seeds,
|
| 568 |
+
inputs=None,
|
| 569 |
+
outputs=[seed_a, seed_b],
|
| 570 |
+
show_progress="hidden",
|
| 571 |
+
)
|
| 572 |
+
generate_button.click(
|
| 573 |
+
generate,
|
| 574 |
+
inputs=[
|
| 575 |
+
model_id,
|
| 576 |
+
scheduler_name,
|
| 577 |
+
prompt_a,
|
| 578 |
+
prompt_b,
|
| 579 |
+
prompt_c,
|
| 580 |
+
negative_prompt,
|
| 581 |
+
embedding_mode,
|
| 582 |
+
prompt_mix,
|
| 583 |
+
analogy_strength,
|
| 584 |
+
renormalize_prompt,
|
| 585 |
+
seed_a,
|
| 586 |
+
seed_b,
|
| 587 |
+
noise_mix,
|
| 588 |
+
renormalize_noise,
|
| 589 |
+
equation_mode,
|
| 590 |
+
guidance_scale,
|
| 591 |
+
prediction_mix,
|
| 592 |
+
negative_coeff,
|
| 593 |
+
prompt_coeff,
|
| 594 |
+
delta_coeff,
|
| 595 |
+
num_steps,
|
| 596 |
+
width,
|
| 597 |
+
height,
|
| 598 |
+
],
|
| 599 |
+
outputs=[output_image, snapshots, summary, metrics, trace],
|
| 600 |
+
show_progress="full",
|
| 601 |
+
)
|
| 602 |
+
|
| 603 |
+
return demo
|
| 604 |
+
|
| 605 |
+
|
| 606 |
+
if __name__ == "__main__":
|
| 607 |
+
build_app().queue(max_size=8).launch()
|
requirements.txt
ADDED
|
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
gradio>=5.22.0
|
| 2 |
+
spaces
|
| 3 |
+
torch>=2.8.0
|
| 4 |
+
diffusers>=0.35.0
|
| 5 |
+
transformers
|
| 6 |
+
accelerate
|
| 7 |
+
safetensors
|
| 8 |
+
pillow
|
| 9 |
+
numpy
|