Add triangle prompt embedding mixer

README.md

@@ -22,6 +22,7 @@ A ZeroGPU-ready Gradio app for learning how Stable Diffusion works inside Diffus
 Instead of only calling `pipe(prompt)`, the app exposes a small custom denoising loop:
 
 - prompt embeddings can be used directly, averaged, or combined with vector arithmetic
+- three prompt embeddings can be explored as a clickable triangle between concepts
 - initial latent noise can come from one seed or a blend of two seeds
 - classifier-free guidance can use standard CFG or a student-edited equation
 - intermediate latent snapshots show how the image emerges across denoising steps

app.py

@@ -40,6 +40,13 @@ prompt_c, _ = encode_prompt(prompt_c, negative_prompt)
 
 if mode == "average":
     prompt_embeds = (1 - mix) * prompt_a + mix * prompt_b
+elif mode == "triangle":
+    total = weight_a + weight_b + weight_c
+    prompt_embeds = (
+        (weight_a / total) * prompt_a
+        + (weight_b / total) * prompt_b
+        + (weight_c / total) * prompt_c
+    )
 elif mode == "analogy":
     prompt_embeds = prompt_a + strength * (prompt_b - prompt_c)
 else:
@@ -120,6 +127,128 @@ def blank_image(message="Run generation to make an image."):
     return image
 
 
+TRIANGLE_SIZE = 360
+TRIANGLE_A = (180, 32)
+TRIANGLE_B = (42, 300)
+TRIANGLE_C = (318, 300)
+
+
+def normalized_triangle_weights(weight_a, weight_b, weight_c):
+    weights = [max(0.0, float(weight_a)), max(0.0, float(weight_b)), max(0.0, float(weight_c))]
+    total = sum(weights)
+    if total <= 0:
+        return 1.0, 0.0, 0.0
+    return tuple(weight / total for weight in weights)
+
+
+def weighted_triangle_point(weight_a, weight_b, weight_c):
+    weight_a, weight_b, weight_c = normalized_triangle_weights(weight_a, weight_b, weight_c)
+    x = weight_a * TRIANGLE_A[0] + weight_b * TRIANGLE_B[0] + weight_c * TRIANGLE_C[0]
+    y = weight_a * TRIANGLE_A[1] + weight_b * TRIANGLE_B[1] + weight_c * TRIANGLE_C[1]
+    return int(round(x)), int(round(y))
+
+
+def short_corner_label(text, fallback):
+    text = " ".join(str(text or fallback).split())
+    return text[:24] + ("..." if len(text) > 24 else "")
+
+
+def make_triangle_picker(prompt_a, prompt_b, prompt_c, weight_a=1 / 3, weight_b=1 / 3, weight_c=1 / 3):
+    image = Image.new("RGB", (TRIANGLE_SIZE, TRIANGLE_SIZE), (248, 250, 252))
+    draw = ImageDraw.Draw(image)
+    vertices = [TRIANGLE_A, TRIANGLE_B, TRIANGLE_C]
+
+    draw.polygon(vertices, fill=(235, 244, 255), outline=(43, 74, 111))
+    for i in range(1, 7):
+        t = i / 7
+        left = (
+            int((1 - t) * TRIANGLE_A[0] + t * TRIANGLE_B[0]),
+            int((1 - t) * TRIANGLE_A[1] + t * TRIANGLE_B[1]),
+        )
+        right = (
+            int((1 - t) * TRIANGLE_A[0] + t * TRIANGLE_C[0]),
+            int((1 - t) * TRIANGLE_A[1] + t * TRIANGLE_C[1]),
+        )
+        draw.line((left, right), fill=(190, 207, 225), width=1)
+
+        bottom = (
+            int((1 - t) * TRIANGLE_B[0] + t * TRIANGLE_C[0]),
+            int((1 - t) * TRIANGLE_B[1] + t * TRIANGLE_C[1]),
+        )
+        left_side = (
+            int((1 - t) * TRIANGLE_A[0] + t * TRIANGLE_B[0]),
+            int((1 - t) * TRIANGLE_A[1] + t * TRIANGLE_B[1]),
+        )
+        right_side = (
+            int((1 - t) * TRIANGLE_A[0] + t * TRIANGLE_C[0]),
+            int((1 - t) * TRIANGLE_A[1] + t * TRIANGLE_C[1]),
+        )
+        draw.line((TRIANGLE_B, right_side), fill=(214, 224, 236), width=1)
+        draw.line((TRIANGLE_C, left_side), fill=(214, 224, 236), width=1)
+        draw.line((bottom, TRIANGLE_A), fill=(214, 224, 236), width=1)
+
+    labels = [
+        (TRIANGLE_A, "A", short_corner_label(prompt_a, "Prompt A"), (55, 94, 151)),
+        (TRIANGLE_B, "B", short_corner_label(prompt_b, "Prompt B"), (5, 122, 85)),
+        (TRIANGLE_C, "C", short_corner_label(prompt_c, "Prompt C"), (154, 72, 174)),
+    ]
+    for (x, y), letter, label, color in labels:
+        draw.ellipse((x - 13, y - 13, x + 13, y + 13), fill=color, outline=(255, 255, 255), width=3)
+        draw.text((x - 4, y - 7), letter, fill=(255, 255, 255))
+        label_x = max(8, min(TRIANGLE_SIZE - 150, x - 70))
+        label_y = y - 34 if y < TRIANGLE_SIZE / 2 else y + 18
+        draw.text((label_x, label_y), label, fill=(30, 41, 59))
+
+    x, y = weighted_triangle_point(weight_a, weight_b, weight_c)
+    draw.ellipse((x - 9, y - 9, x + 9, y + 9), fill=(239, 68, 68), outline=(15, 23, 42), width=2)
+    draw.text((12, 12), "Click inside the triangle to choose the embedding blend.", fill=(51, 65, 85))
+    return image
+
+
+def barycentric_triangle_weights(x, y):
+    ax, ay = TRIANGLE_A
+    bx, by = TRIANGLE_B
+    cx, cy = TRIANGLE_C
+    denominator = (by - cy) * (ax - cx) + (cx - bx) * (ay - cy)
+    if denominator == 0:
+        return 1.0, 0.0, 0.0
+
+    weight_a = ((by - cy) * (x - cx) + (cx - bx) * (y - cy)) / denominator
+    weight_b = ((cy - ay) * (x - cx) + (ax - cx) * (y - cy)) / denominator
+    weight_c = 1.0 - weight_a - weight_b
+
+    if min(weight_a, weight_b, weight_c) < 0:
+        weight_a, weight_b, weight_c = normalized_triangle_weights(weight_a, weight_b, weight_c)
+
+    return normalized_triangle_weights(weight_a, weight_b, weight_c)
+
+
+def triangle_status(weight_a, weight_b, weight_c):
+    weight_a, weight_b, weight_c = normalized_triangle_weights(weight_a, weight_b, weight_c)
+    return f"A: {weight_a:.2f}   B: {weight_b:.2f}   C: {weight_c:.2f}"
+
+
+def update_triangle_from_weights(prompt_a, prompt_b, prompt_c, weight_a, weight_b, weight_c):
+    weight_a, weight_b, weight_c = normalized_triangle_weights(weight_a, weight_b, weight_c)
+    return (
+        make_triangle_picker(prompt_a, prompt_b, prompt_c, weight_a, weight_b, weight_c),
+        triangle_status(weight_a, weight_b, weight_c),
+        weight_a,
+        weight_b,
+        weight_c,
+    )
+
+
+def select_triangle_point(prompt_a, prompt_b, prompt_c, evt: gr.SelectData):
+    index = evt.index
+    if isinstance(index, dict):
+        x, y = index.get("x", TRIANGLE_A[0]), index.get("y", TRIANGLE_A[1])
+    else:
+        x, y = index[:2]
+    weight_a, weight_b, weight_c = barycentric_triangle_weights(float(x), float(y))
+    return update_triangle_from_weights(prompt_a, prompt_b, prompt_c, weight_a, weight_b, weight_c)
+
+
 @lru_cache(maxsize=2)
 def load_pipe(model_id, scheduler_name, device_type):
     device = torch.device(device_type)
@@ -185,6 +314,9 @@ def mix_prompt_embeddings(
     negative_prompt,
     embedding_mode,
     prompt_mix,
+    triangle_weight_a,
+    triangle_weight_b,
+    triangle_weight_c,
     analogy_strength,
     renormalize_prompt,
 ):
@@ -195,6 +327,17 @@ def mix_prompt_embeddings(
     if embedding_mode == "Average prompt A and B":
         mixed = (1.0 - prompt_mix) * emb_a + prompt_mix * emb_b
         formula = f"prompt = {(1.0 - prompt_mix):.2f} * A + {prompt_mix:.2f} * B"
+    elif embedding_mode == "Triangle blend: A/B/C":
+        triangle_weight_a, triangle_weight_b, triangle_weight_c = normalized_triangle_weights(
+            triangle_weight_a,
+            triangle_weight_b,
+            triangle_weight_c,
+        )
+        mixed = triangle_weight_a * emb_a + triangle_weight_b * emb_b + triangle_weight_c * emb_c
+        formula = (
+            f"prompt = {triangle_weight_a:.2f} * A + {triangle_weight_b:.2f} * B "
+            f"+ {triangle_weight_c:.2f} * C"
+        )
     elif embedding_mode == "Vector arithmetic: A + s * (B - C)":
         mixed = emb_a + analogy_strength * (emb_b - emb_c)
         formula = f"prompt = A + {analogy_strength:.2f} * (B - C)"
@@ -212,6 +355,7 @@ def mix_prompt_embeddings(
         ["cosine(A, B)", round(cosine_similarity(emb_a, emb_b), 4)],
         ["cosine(A, mixed)", round(cosine_similarity(emb_a, mixed), 4)],
         ["cosine(B, mixed)", round(cosine_similarity(emb_b, mixed), 4)],
+        ["cosine(C, mixed)", round(cosine_similarity(emb_c, mixed), 4)],
         ["norm(A)", round(float(original_norm.cpu()), 3)],
         ["norm(mixed)", round(float(mixed.detach().float().norm().cpu()), 3)],
     ]
@@ -316,6 +460,9 @@ def generate(
     negative_prompt,
     embedding_mode,
     prompt_mix,
+    triangle_weight_a,
+    triangle_weight_b,
+    triangle_weight_c,
     analogy_strength,
     renormalize_prompt,
     seed_a,
@@ -356,6 +503,9 @@ def generate(
         negative_prompt or "",
         embedding_mode,
         float(prompt_mix),
+        float(triangle_weight_a),
+        float(triangle_weight_b),
+        float(triangle_weight_c),
         float(analogy_strength),
         bool(renormalize_prompt),
     )
@@ -486,12 +636,28 @@ def build_app():
                         [
                             "Prompt A only",
                             "Average prompt A and B",
+                            "Triangle blend: A/B/C",
                             "Vector arithmetic: A + s * (B - C)",
                         ],
-                        value="Average prompt A and B",
+                        value="Triangle blend: A/B/C",
                         label="Embedding equation",
                     )
                     prompt_mix = gr.Slider(0, 1, value=0.5, step=0.05, label="Prompt B weight")
+                    triangle_picker = gr.Image(
+                        value=make_triangle_picker(DEFAULT_PROMPT_A, DEFAULT_PROMPT_B, DEFAULT_PROMPT_C),
+                        label="Triangle embedding mixer",
+                        type="pil",
+                        interactive=False,
+                    )
+                    triangle_status_box = gr.Textbox(
+                        value=triangle_status(1 / 3, 1 / 3, 1 / 3),
+                        label="Triangle weights",
+                        interactive=False,
+                    )
+                    with gr.Row():
+                        triangle_weight_a = gr.Slider(0, 1, value=1 / 3, step=0.01, label="A weight")
+                        triangle_weight_b = gr.Slider(0, 1, value=1 / 3, step=0.01, label="B weight")
+                        triangle_weight_c = gr.Slider(0, 1, value=1 / 3, step=0.01, label="C weight")
                     analogy_strength = gr.Slider(-2, 2, value=0.8, step=0.1, label="Vector arithmetic strength")
                     renormalize_prompt = gr.Checkbox(value=True, label="Keep mixed prompt embedding norm near prompt A")
 
@@ -569,6 +735,19 @@ def build_app():
             outputs=[seed_a, seed_b],
             show_progress="hidden",
         )
+        triangle_picker.select(
+            select_triangle_point,
+            inputs=[prompt_a, prompt_b, prompt_c],
+            outputs=[triangle_picker, triangle_status_box, triangle_weight_a, triangle_weight_b, triangle_weight_c],
+            show_progress="hidden",
+        )
+        for triangle_input in [prompt_a, prompt_b, prompt_c, triangle_weight_a, triangle_weight_b, triangle_weight_c]:
+            triangle_input.change(
+                update_triangle_from_weights,
+                inputs=[prompt_a, prompt_b, prompt_c, triangle_weight_a, triangle_weight_b, triangle_weight_c],
+                outputs=[triangle_picker, triangle_status_box, triangle_weight_a, triangle_weight_b, triangle_weight_c],
+                show_progress="hidden",
+            )
         generate_button.click(
             generate,
             inputs=[
@@ -580,6 +759,9 @@ def build_app():
                 negative_prompt,
                 embedding_mode,
                 prompt_mix,
+                triangle_weight_a,
+                triangle_weight_b,
+                triangle_weight_c,
                 analogy_strength,
                 renormalize_prompt,
                 seed_a,