Blame for app.py

"""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
# ----------------------------------------------------------------------------


@lru_cache(maxsize=256)
@torch.no_grad()
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)


@torch.no_grad()
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
# ----------------------------------------------------------------------------


@gpu
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
# ----------------------------------------------------------------------------


@gpu
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


@gpu
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


@gpu
@torch.no_grad()
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()