Leeps commited on
Commit
c37d50d
·
1 Parent(s): b9cdbb1

Add Stable Diffusion equation playground

Browse files
Files changed (3) hide show
  1. README.md +41 -7
  2. app.py +607 -0
  3. requirements.txt +9 -0
README.md CHANGED
@@ -1,13 +1,47 @@
1
  ---
2
- title: Diffusers Playground
3
- emoji: 🐨
4
- colorFrom: purple
5
- colorTo: blue
6
  sdk: gradio
7
- sdk_version: 6.19.0
8
- python_version: '3.12'
9
  app_file: app.py
10
  pinned: false
 
 
 
 
 
11
  ---
12
 
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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