Spaces:
Paused
Paused
| """Text Embedding Explorer for Stable Diffusion 1.5. | |
| A Gradio Space for BA creative coding students. Four tabs let you poke at the | |
| text embedding (the 77 x 768 grid of numbers) that Stable Diffusion 1.5 gets | |
| from its CLIP text encoder, instead of at the prompt text itself. | |
| Tabs: | |
| 1. Interpolate blend two prompts in embedding space | |
| 2. Arithmetic A - B + C on embeddings | |
| 3. Token editing scale or blend single tokens | |
| 4. Map 2D map (PCA) and similarity matrix of embeddings | |
| """ | |
| import os | |
| import random | |
| from functools import lru_cache | |
| import gradio as gr | |
| import numpy as np | |
| import pandas as pd | |
| import plotly.graph_objects as go | |
| import torch | |
| from diffusers import StableDiffusionPipeline, UniPCMultistepScheduler | |
| from sklearn.decomposition import PCA | |
| # On ZeroGPU Spaces the `spaces` package provides a GPU decorator. | |
| # On a dedicated GPU (or locally) we simply do nothing. | |
| try: | |
| import spaces | |
| gpu = spaces.GPU | |
| except ImportError: | |
| def gpu(fn): | |
| return fn | |
| # ---------------------------------------------------------------------------- | |
| # Model | |
| # ---------------------------------------------------------------------------- | |
| MODEL_ID = os.environ.get("MODEL_ID", "stable-diffusion-v1-5/stable-diffusion-v1-5") | |
| # Set DISABLE_SAFETY_CHECKER=1 in the Space settings to switch the NSFW filter off. | |
| # Odd embedding mixes can trigger false positives, which show up as black images. | |
| USE_SAFETY_CHECKER = os.environ.get("DISABLE_SAFETY_CHECKER", "0") != "1" | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 | |
| _pipe_kwargs = dict(dtype=DTYPE, variant="fp16", use_safetensors=True) | |
| if not USE_SAFETY_CHECKER: | |
| _pipe_kwargs.update(safety_checker=None, requires_safety_checker=False) | |
| pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, **_pipe_kwargs) | |
| pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config) | |
| pipe = pipe.to(DEVICE) | |
| pipe.set_progress_bar_config(disable=True) | |
| tokenizer = pipe.tokenizer | |
| text_encoder = pipe.text_encoder | |
| MAX_LEN = tokenizer.model_max_length # 77 for SD 1.5 | |
| # ---------------------------------------------------------------------------- | |
| # Embedding helpers | |
| # ---------------------------------------------------------------------------- | |
| def encode(text: str) -> torch.Tensor: | |
| """Prompt -> text embedding of shape (77, 768), float32, on DEVICE. | |
| This is exactly what the UNet of SD 1.5 receives as its conditioning. | |
| Positions beyond the prompt are filled with the end-of-text token. | |
| """ | |
| ids = tokenizer( | |
| text, | |
| padding="max_length", | |
| max_length=MAX_LEN, | |
| truncation=True, | |
| return_tensors="pt", | |
| ).input_ids.to(DEVICE) | |
| return text_encoder(ids)[0][0].float() | |
| def token_strings(text: str): | |
| """Readable tokens of the prompt (without start and end token).""" | |
| ids = tokenizer(text, truncation=True, max_length=MAX_LEN).input_ids | |
| toks = tokenizer.convert_ids_to_tokens(ids)[1:-1] | |
| return [t.replace("</w>", "") for t in toks] | |
| def is_truncated(text: str) -> bool: | |
| return len(tokenizer(text).input_ids) > MAX_LEN | |
| def slerp(a: torch.Tensor, b: torch.Tensor, t: float) -> torch.Tensor: | |
| """Spherical interpolation, applied to every token vector separately. | |
| Direction goes along the arc between the two unit vectors, the length | |
| changes linearly. Works for t outside 0..1 as well (extrapolation). | |
| """ | |
| na = a.norm(dim=-1, keepdim=True) | |
| nb = b.norm(dim=-1, keepdim=True) | |
| ua, ub = a / na, b / nb | |
| dot = (ua * ub).sum(-1, keepdim=True).clamp(-1.0, 1.0) | |
| omega = torch.acos(dot) | |
| so = torch.sin(omega) | |
| direction = (torch.sin((1 - t) * omega) * ua + torch.sin(t * omega) * ub) / so.clamp_min(1e-6) | |
| out = direction * ((1 - t) * na + t * nb) | |
| # Where the two vectors point the same way, slerp is undefined. Use lerp there. | |
| return torch.where(so < 1e-4, (1 - t) * a + t * b, out) | |
| def lerp(a: torch.Tensor, b: torch.Tensor, t: float) -> torch.Tensor: | |
| return (1 - t) * a + t * b | |
| def cos_sim(a: torch.Tensor, b: torch.Tensor) -> float: | |
| return torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() | |
| def parse_seed(seed) -> int: | |
| try: | |
| s = int(seed) | |
| except (TypeError, ValueError): | |
| s = -1 | |
| return s if s >= 0 else random.randint(0, 2**32 - 1) | |
| def render(embeds, negative_text, steps, cfg, seed, size, progress=None): | |
| """Generate one image per row of `embeds` (N, 77, 768), all from the same noise.""" | |
| n = embeds.shape[0] | |
| neg = encode(negative_text or "") | |
| g = torch.Generator(device="cpu").manual_seed(seed) | |
| latents = torch.randn((1, 4, size // 8, size // 8), generator=g).to(DEVICE, DTYPE) | |
| images = [] | |
| chunk = 4 | |
| for start in range(0, n, chunk): | |
| part = embeds[start : start + chunk] | |
| k = part.shape[0] | |
| if progress is not None: | |
| progress(start / n, desc=f"Generating image {start + 1} of {n}") | |
| out = pipe( | |
| prompt_embeds=part.to(DTYPE), | |
| negative_prompt_embeds=neg.to(DTYPE).unsqueeze(0).expand(k, -1, -1), | |
| latents=latents.expand(k, -1, -1, -1).contiguous(), | |
| num_inference_steps=int(steps), | |
| guidance_scale=float(cfg), | |
| height=int(size), | |
| width=int(size), | |
| ) | |
| images.extend(out.images) | |
| return images | |
| def need_text(*texts): | |
| for t in texts: | |
| if not t or not t.strip(): | |
| raise gr.Error("Please fill in all prompt fields of this tab.") | |
| # ---------------------------------------------------------------------------- | |
| # Tab 1: Interpolation | |
| # ---------------------------------------------------------------------------- | |
| def run_interpolation(a_text, b_text, frames, method, t_min, t_max, | |
| neg, steps, cfg, seed, size, progress=gr.Progress()): | |
| need_text(a_text, b_text) | |
| seed = parse_seed(seed) | |
| a, b = encode(a_text), encode(b_text) | |
| blend = slerp if method == "slerp" else lerp | |
| ts = np.linspace(float(t_min), float(t_max), int(frames)) | |
| embeds = torch.stack([blend(a, b, float(t)) for t in ts]) | |
| images = render(embeds, neg, steps, cfg, seed, int(size), progress) | |
| captions = [f"t = {t:.2f}" for t in ts] | |
| info = ( | |
| f"Seed {seed}. Similarity between the embeddings of A and B (cosine, 1 = identical): " | |
| f"{cos_sim(a, b):.3f}." | |
| ) | |
| return list(zip(images, captions)), info | |
| # ---------------------------------------------------------------------------- | |
| # Tab 2: Arithmetic | |
| # ---------------------------------------------------------------------------- | |
| def run_arithmetic(a_text, b_text, c_text, w_b, w_c, renorm, show_inputs, | |
| neg, steps, cfg, seed, size, progress=gr.Progress()): | |
| need_text(a_text, b_text, c_text) | |
| seed = parse_seed(seed) | |
| a, b, c = encode(a_text), encode(b_text), encode(c_text) | |
| result = a - float(w_b) * b + float(w_c) * c | |
| if renorm: | |
| # Give every token vector the same length it had in A. | |
| result = result / result.norm(dim=-1, keepdim=True).clamp_min(1e-6) * a.norm(dim=-1, keepdim=True) | |
| result[0] = a[0] # the start token is identical for every prompt, keep it as is | |
| label = f"A - {w_b:g}·B + {w_c:g}·C" | |
| if show_inputs: | |
| embeds = torch.stack([a, b, c, result]) | |
| captions = ["A", "B", "C", label] | |
| else: | |
| embeds = result.unsqueeze(0) | |
| captions = [label] | |
| images = render(embeds, neg, steps, cfg, seed, int(size), progress) | |
| info = f"Seed {seed}. Cosine similarity of the result to A: {cos_sim(result, a):.3f}." | |
| return list(zip(images, captions)), info | |
| # ---------------------------------------------------------------------------- | |
| # Tab 3: Token editing | |
| # ---------------------------------------------------------------------------- | |
| TOKEN_COLUMNS = ["pos", "token", "scale", "blend_with_B"] | |
| def make_token_table(a_text): | |
| if not a_text or not a_text.strip(): | |
| raise gr.Error("Write prompt A first.") | |
| toks = token_strings(a_text) | |
| df = pd.DataFrame( | |
| { | |
| "pos": list(range(1, len(toks) + 1)), | |
| "token": toks, | |
| "scale": [1.0] * len(toks), | |
| "blend_with_B": [0.0] * len(toks), | |
| } | |
| ) | |
| msg = f"{len(toks)} tokens." | |
| if is_truncated(a_text): | |
| msg += f" The prompt was cut off at {MAX_LEN - 2} tokens, the model ignores the rest." | |
| return df, msg | |
| def run_token_edit(a_text, b_text, table, neg, steps, cfg, seed, size, progress=gr.Progress()): | |
| need_text(a_text) | |
| seed = parse_seed(seed) | |
| current = token_strings(a_text) | |
| if table is None or len(table) != len(current) or list(table["token"]) != current: | |
| raise gr.Error("The prompt changed since the table was made. Click 'Split into tokens' again.") | |
| a = encode(a_text) | |
| b = encode(b_text or "") # empty B is allowed, it then means 'nothing' | |
| empty = encode("") | |
| edited = a.clone() | |
| changed = 0 | |
| for _, row in table.iterrows(): | |
| pos = int(row["pos"]) | |
| scale = float(np.nan_to_num(row["scale"], nan=1.0)) | |
| mix = float(np.clip(np.nan_to_num(row["blend_with_B"], nan=0.0), 0.0, 1.0)) | |
| vec = (1 - mix) * a[pos] + mix * b[pos] | |
| # Scale is measured from the embedding of the empty prompt: | |
| # 1 = unchanged, 0 = "no information", 2 = twice as far from "no information". | |
| vec = empty[pos] + scale * (vec - empty[pos]) | |
| if scale != 1.0 or mix != 0.0: | |
| changed += 1 | |
| edited[pos] = vec | |
| images = render(torch.stack([a, edited]), neg, steps, cfg, seed, int(size), progress) | |
| captions = ["original", "edited"] | |
| info = f"Seed {seed}. {changed} token(s) edited. Cosine similarity edited vs original: {cos_sim(a, edited):.3f}." | |
| return list(zip(images, captions)), info | |
| # ---------------------------------------------------------------------------- | |
| # Tab 4: Map of embeddings | |
| # ---------------------------------------------------------------------------- | |
| POINT_COLOR = "#3987e5" # blue that stays readable on light and dark backgrounds | |
| INK = "#8a8a85" # mid gray for axes and labels, works on both backgrounds | |
| MAX_MAP_ENTRIES = 60 | |
| def run_map(entries_text, pooling): | |
| labels = [] | |
| for line in (entries_text or "").splitlines(): | |
| line = line.strip() | |
| if line and line not in labels: | |
| labels.append(line) | |
| if len(labels) < 3: | |
| raise gr.Error("Please give at least 3 different entries, one per line.") | |
| if len(labels) > MAX_MAP_ENTRIES: | |
| raise gr.Error(f"Please use at most {MAX_MAP_ENTRIES} entries.") | |
| vecs = [] | |
| for text in labels: | |
| emb = encode(text) | |
| n = min(len(token_strings(text)), MAX_LEN - 2) | |
| if pooling.startswith("Average"): | |
| vecs.append(emb[1 : n + 1].mean(0)) | |
| else: | |
| vecs.append(emb[n + 1]) # first end-of-text token | |
| x = torch.stack(vecs).cpu().numpy() | |
| pca = PCA(n_components=2) | |
| xy = pca.fit_transform(x) | |
| var = pca.explained_variance_ratio_ * 100 | |
| scatter = go.Figure( | |
| go.Scatter( | |
| x=xy[:, 0], | |
| y=xy[:, 1], | |
| mode="markers+text", | |
| text=labels, | |
| textposition="top center", | |
| textfont=dict(color=INK, size=12), | |
| marker=dict(size=10, color=POINT_COLOR, line=dict(width=2, color="rgba(255,255,255,0.6)")), | |
| hovertemplate="%{text}<extra></extra>", | |
| ) | |
| ) | |
| scatter.update_layout( | |
| title="Embedding map (PCA, 2 of 768 dimensions)", | |
| xaxis=dict(title=f"PC1 ({var[0]:.0f}% of variance)", showgrid=False, zeroline=False, color=INK), | |
| yaxis=dict(title=f"PC2 ({var[1]:.0f}% of variance)", showgrid=False, zeroline=False, color=INK), | |
| paper_bgcolor="rgba(0,0,0,0)", | |
| plot_bgcolor="rgba(0,0,0,0)", | |
| font=dict(color=INK), | |
| margin=dict(l=40, r=20, t=50, b=40), | |
| height=520, | |
| showlegend=False, | |
| ) | |
| xn = x / np.linalg.norm(x, axis=1, keepdims=True) | |
| sim = xn @ xn.T | |
| off_diag = sim[~np.eye(len(labels), dtype=bool)] | |
| heat = go.Figure( | |
| go.Heatmap( | |
| z=sim, | |
| x=labels, | |
| y=labels, | |
| zmin=float(off_diag.min()), | |
| zmax=1.0, | |
| colorscale=[[0, "#eaf2fc"], [1, "#1f5fb0"]], | |
| colorbar=dict(title="cosine"), | |
| hovertemplate="%{y} / %{x}: %{z:.2f}<extra></extra>", | |
| ) | |
| ) | |
| heat.update_layout( | |
| title="Cosine similarity (darker = more similar)", | |
| xaxis=dict(color=INK, tickangle=-45), | |
| yaxis=dict(color=INK, autorange="reversed"), | |
| paper_bgcolor="rgba(0,0,0,0)", | |
| plot_bgcolor="rgba(0,0,0,0)", | |
| font=dict(color=INK), | |
| margin=dict(l=40, r=20, t=50, b=40), | |
| height=560, | |
| ) | |
| note = ( | |
| f"{len(labels)} entries. The two axes together keep {var.sum():.0f}% of the variation " | |
| "of the original 768 numbers, so distances on the map are only a rough picture." | |
| ) | |
| return scatter, heat, note | |
| # ---------------------------------------------------------------------------- | |
| # Interface | |
| # ---------------------------------------------------------------------------- | |
| INTRO = """ | |
| # Text Embedding Explorer for Stable Diffusion 1.5 | |
| Stable Diffusion never sees your prompt as text. A CLIP text encoder first turns it into a grid of | |
| numbers, the **text embedding**, and only this grid steers the image. Here you can change the grid | |
| directly and see what happens. Every tab uses the **same random noise (seed)** for all its images, so | |
| the differences you see come from the embedding and nothing else. | |
| """ | |
| EXPLAIN = """ | |
| **From prompt to numbers.** The prompt is cut into *tokens* (word pieces, at most 75). Each token becomes | |
| a vector of 768 numbers. The model always works with 77 slots: a start token, your tokens, an end token | |
| and then more end tokens as padding. So the embedding is a table of **77 rows x 768 columns**. | |
| **Why does editing it work?** Similar meanings end up as similar numbers. If you move the numbers a little | |
| toward another prompt, the image moves a little toward that prompt too. | |
| **Good to know.** Because a token sees the tokens before it, the vector of a word depends a bit on the words | |
| in front of it. Embeddings of prompts with different lengths do not line up perfectly, row by row. | |
| That is part of the fun and part of the glitchiness. | |
| """ | |
| with gr.Blocks(title="Text Embedding Explorer for Stable Diffusion 1.5") as demo: | |
| gr.Markdown(INTRO) | |
| with gr.Accordion("What is a text embedding?", open=False): | |
| gr.Markdown(EXPLAIN) | |
| with gr.Accordion("Generation settings (shared by all tabs)", open=False): | |
| with gr.Row(): | |
| steps = gr.Slider(8, 48, value=16, step=1, label="Steps") | |
| cfg = gr.Slider(1, 15, value=7.5, step=0.5, label="Guidance scale") | |
| size = gr.Dropdown([384, 448, 512], value=512, label="Image size (px)") | |
| with gr.Row(): | |
| seed = gr.Number(value=42, precision=0, label="Seed (negative = random)") | |
| rand_btn = gr.Button("Randomise seed", scale=0) | |
| neg = gr.Textbox(label="Negative prompt (optional)", placeholder="e.g. blurry, low quality") | |
| rand_btn.click(lambda: random.randint(0, 2**32 - 1), outputs=seed) | |
| shared = [neg, steps, cfg, seed, size] | |
| with gr.Tabs(): | |
| # -------------------------------------------------- Interpolate | |
| with gr.Tab("1. Interpolate"): | |
| gr.Markdown( | |
| "Blend prompt **A** into prompt **B**. At t = 0 you get A, at t = 1 you get B, " | |
| "in between the embeddings are mixed." | |
| ) | |
| with gr.Row(): | |
| i_a = gr.Textbox(label="Prompt A", value="a watercolor painting of a red fox in a forest") | |
| i_b = gr.Textbox(label="Prompt B", value="a watercolor painting of a blue whale in the ocean") | |
| with gr.Row(): | |
| i_frames = gr.Slider(2, 9, value=5, step=1, label="Number of images") | |
| i_method = gr.Radio( | |
| ["slerp", "lerp"], value="slerp", label="Blend method", | |
| info="slerp follows an arc, lerp a straight line. For prompts they look alike.", | |
| ) | |
| with gr.Accordion("Go beyond A and B", open=False): | |
| gr.Markdown("Values below 0 or above 1 push past the prompts (extrapolation). Expect strange results.") | |
| with gr.Row(): | |
| i_tmin = gr.Slider(-1, 1, value=0, step=0.05, label="First t") | |
| i_tmax = gr.Slider(0, 2, value=1, step=0.05, label="Last t") | |
| i_btn = gr.Button("Generate", variant="primary") | |
| i_gal = gr.Gallery(label="Result", columns=5, height="auto", object_fit="contain") | |
| i_info = gr.Markdown() | |
| i_btn.click( | |
| run_interpolation, | |
| [i_a, i_b, i_frames, i_method, i_tmin, i_tmax, *shared], | |
| [i_gal, i_info], | |
| ) | |
| gr.Examples( | |
| [ | |
| ["a photo of a cat sleeping on a sofa", "a photo of a dog sleeping on a sofa"], | |
| ["a city street in summer, golden hour", "a city street in winter, snowing, night"], | |
| ], | |
| [i_a, i_b], | |
| cache_examples=False, | |
| ) | |
| # -------------------------------------------------- Arithmetic | |
| with gr.Tab("2. Arithmetic"): | |
| gr.Markdown( | |
| "Add and subtract embeddings: **A minus B plus C**. The classic idea is *king - man + woman*. " | |
| "It works best when the three prompts have the same sentence structure and length." | |
| ) | |
| with gr.Row(): | |
| a_a = gr.Textbox(label="A (start)", value="a photo of a king") | |
| a_b = gr.Textbox(label="B (subtract)", value="a photo of a man") | |
| a_c = gr.Textbox(label="C (add)", value="a photo of a woman") | |
| with gr.Row(): | |
| a_wb = gr.Slider(0, 2, value=1, step=0.05, label="Weight of B") | |
| a_wc = gr.Slider(0, 2, value=1, step=0.05, label="Weight of C") | |
| with gr.Row(): | |
| a_renorm = gr.Checkbox( | |
| value=True, label="Keep vector lengths like A", | |
| info="Rescales each token vector to the length it has in A. Usually keeps images cleaner.", | |
| ) | |
| a_show = gr.Checkbox(value=True, label="Also show A, B and C") | |
| a_btn = gr.Button("Generate", variant="primary") | |
| a_gal = gr.Gallery(label="Result", columns=4, height="auto", object_fit="contain") | |
| a_info = gr.Markdown() | |
| a_btn.click( | |
| run_arithmetic, | |
| [a_a, a_b, a_c, a_wb, a_wc, a_renorm, a_show, *shared], | |
| [a_gal, a_info], | |
| ) | |
| gr.Examples( | |
| [ | |
| ["a photo of a king", "a photo of a man", "a photo of a woman"], | |
| ["a painting of a summer meadow", "a painting of summer", "a painting of winter"], | |
| ["a photo of a kitten", "a photo of a cat", "a photo of a dog"], | |
| ], | |
| [a_a, a_b, a_c], | |
| cache_examples=False, | |
| ) | |
| # -------------------------------------------------- Token editing | |
| with gr.Tab("3. Edit tokens"): | |
| gr.Markdown( | |
| "1. Write prompt **A** and click **Split into tokens**.\n" | |
| "2. Edit the table. **scale**: 1 = unchanged, 0 = token is muted, 2 = token is exaggerated. " | |
| "**blend_with_B**: 0 = keep A, 1 = take this slot from prompt B (optional).\n" | |
| "3. Click **Generate** to see the original next to your edit." | |
| ) | |
| with gr.Row(): | |
| t_a = gr.Textbox(label="Prompt A", value="a red balloon floating over a snowy mountain, oil painting") | |
| t_b = gr.Textbox(label="Prompt B (optional)", placeholder="Slots are taken from B at the same position") | |
| t_split = gr.Button("Split into tokens") | |
| t_msg = gr.Markdown() | |
| t_table = gr.Dataframe( | |
| headers=TOKEN_COLUMNS, | |
| datatype=["number", "str", "number", "number"], | |
| column_count=(4, "fixed"), | |
| static_columns=[0, 1], | |
| interactive=True, | |
| label="Tokens (edit the last two columns)", | |
| ) | |
| t_btn = gr.Button("Generate", variant="primary") | |
| t_gal = gr.Gallery(label="Result", columns=2, height="auto", object_fit="contain") | |
| t_info = gr.Markdown() | |
| t_split.click(make_token_table, t_a, [t_table, t_msg]) | |
| t_a.submit(make_token_table, t_a, [t_table, t_msg]) | |
| t_btn.click(run_token_edit, [t_a, t_b, t_table, *shared], [t_gal, t_info]) | |
| demo.load(make_token_table, t_a, [t_table, t_msg]) | |
| # -------------------------------------------------- Map | |
| with gr.Tab("4. Map"): | |
| gr.Markdown( | |
| "Type words or short prompts, one per line. Each becomes a point. " | |
| "Close points have similar embeddings. This is the same space the other tabs move through." | |
| ) | |
| m_text = gr.Textbox( | |
| label="Entries (one per line, 3 to 60)", | |
| lines=10, | |
| value="cat\ndog\nhorse\nlion\ncar\nbicycle\nairplane\nboat\nred\nblue\ngreen\nyellow\nhappy\nsad\nangry", | |
| ) | |
| m_pool = gr.Radio( | |
| ["Average of its tokens", "End token (the model's summary of the whole prompt)"], | |
| value="Average of its tokens", | |
| label="How to turn an entry into one point", | |
| ) | |
| m_btn = gr.Button("Draw map", variant="primary") | |
| m_note = gr.Markdown() | |
| m_scatter = gr.Plot(label="PCA map") | |
| m_heat = gr.Plot(label="Similarity") | |
| m_btn.click(run_map, [m_text, m_pool], [m_scatter, m_heat, m_note]) | |
| gr.Markdown( | |
| "Model: Stable Diffusion 1.5 (`stable-diffusion-v1-5/stable-diffusion-v1-5`), " | |
| "used under its CreativeML OpenRAIL-M license. Built with diffusers and Gradio." | |
| ) | |
| demo.queue(max_size=20) | |
| if __name__ == "__main__": | |
| demo.launch() | |