Update app.py

app.py

@@ -1,69 +1,57 @@
-# app.py ────────────────────────────────────────────────────────────────────
-import io, os, json, math, random, warnings, gc, functools, hashlib
+# app.py ────────────────────────────────────────────────────────────────
+import io, warnings, numpy as np, matplotlib.pyplot as plt
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple
 
 import gradio as gr
-import numpy as np
-import matplotlib.pyplot as plt
+import torch, torch.nn.functional as F
 from PIL import Image
-
-import torch
-import torch.nn.functional as F
 from transformers import T5Tokenizer, T5EncoderModel
 from diffusers import (
     StableDiffusionXLPipeline,
-    DDIMScheduler,
-    EulerDiscreteScheduler,
-    DPMSolverMultistepScheduler,
+    DDIMScheduler, EulerDiscreteScheduler, DPMSolverMultistepScheduler,
 )
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
-# -------------------------------------------------------------------------
 # local modules
 from two_stream_shunt_adapter import TwoStreamShuntAdapter
-from .conditioning_shifter import ConditioningShifter
+from conditioning_shifter import ConditioningShifter, ShiftConfig, AdapterOutput
+from embedding_manager import get_bank
 from configs import T5_SHUNT_REPOS
-from embedding_manager import get_bank   # ← NEW
 
 warnings.filterwarnings("ignore")
 
+# ─── GLOBALS ────────────────────────────────────────────────────────────
+dtype  = torch.float16
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# ───────────────────────────────────────────────────────────────────────────
-# GLOBALS
-# ───────────────────────────────────────────────────────────────────────────
-dtype   = torch.float16
-device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-bank    = get_bank()                           # shared singleton
+_bank  = get_bank()                      # singleton – optional caching
 
-_t5_tok: Optional[T5Tokenizer]    = None
-_t5_mod: Optional[T5EncoderModel] = None
-_pipe:   Optional[StableDiffusionXLPipeline] = None
+_t5_tok: Optional[T5Tokenizer]     = None
+_t5_mod: Optional[T5EncoderModel]  = None
+_pipe  : Optional[StableDiffusionXLPipeline] = None
 
 SCHEDULERS = {
-    "DPM++ 2M":  DPMSolverMultistepScheduler,
-    "DDIM":      DDIMScheduler,
-    "Euler":     EulerDiscreteScheduler,
+    "DPM++ 2M": DPMSolverMultistepScheduler,
+    "DDIM":     DDIMScheduler,
+    "Euler":    EulerDiscreteScheduler,
 }
 
-# easy access to adapter repo metadata
+# adapter-meta from configs.py
 clip_l_opts = T5_SHUNT_REPOS["clip_l"]["shunts_available"]["shunt_list"]
 clip_g_opts = T5_SHUNT_REPOS["clip_g"]["shunts_available"]["shunt_list"]
-repo_l      = T5_SHUNT_REPOS["clip_l"]["repo"]
-repo_g      = T5_SHUNT_REPOS["clip_g"]["repo"]
-conf_l      = T5_SHUNT_REPOS["clip_l"]["config"]
-conf_g      = T5_SHUNT_REPOS["clip_g"]["config"]
+repo_l, conf_l = T5_SHUNT_REPOS["clip_l"]["repo"], T5_SHUNT_REPOS["clip_l"]["config"]
+repo_g, conf_g = T5_SHUNT_REPOS["clip_g"]["repo"], T5_SHUNT_REPOS["clip_g"]["config"]
 
 
-# ───────────────────────────────────────────────────────────────────────────
-# HELPERs
-# ───────────────────────────────────────────────────────────────────────────
+# ─── INITIALISERS ────────────────────────────────────────────────────────
 def _init_t5():
     global _t5_tok, _t5_mod
     if _t5_tok is None:
         _t5_tok = T5Tokenizer.from_pretrained("google/flan-t5-base")
-        _t5_mod = T5EncoderModel.from_pretrained("google/flan-t5-base").to(device).eval()
+        _t5_mod = T5EncoderModel.from_pretrained("google/flan-t5-base") \
+                                .to(device).eval()
 
 
 def _init_pipe():
@@ -71,124 +59,82 @@ def _init_pipe():
     if _pipe is None:
         _pipe = StableDiffusionXLPipeline.from_pretrained(
             "stabilityai/stable-diffusion-xl-base-1.0",
-            torch_dtype=dtype,
-            use_safetensors=True,
-            variant="fp16",
+            torch_dtype=dtype, variant="fp16", use_safetensors=True
         ).to(device)
         _pipe.enable_xformers_memory_efficient_attention()
 
 
-def load_adapter(repo: str, filename: str, cfg: dict):
-    """load a TwoStreamShuntAdapter from HF Hub safetensors"""
-    path = hf_hub_download(repo_id=repo, filename=filename)
-    model = TwoStreamShuntAdapter(cfg).eval()
-    tensors = load_file(path)
-    model.load_state_dict(tensors)
+# ─── HELPERS ─────────────────────────────────────────────────────────────
+def load_adapter(repo: str, filename: str, cfg: dict,
+                 device: torch.device) -> TwoStreamShuntAdapter:
+    path   = hf_hub_download(repo_id=repo, filename=filename)
+    model  = TwoStreamShuntAdapter(cfg).eval()
+    model.load_state_dict(load_file(path))
     return model.to(device)
 
 
 def plot_heat(mat: torch.Tensor | np.ndarray, title: str) -> np.ndarray:
     if isinstance(mat, torch.Tensor):
         mat = mat.detach().cpu().numpy()
-
     if mat.ndim == 1:
         mat = mat[None, :]
-    elif mat.ndim >= 3:                 # (B,T,D) → mean over B
+    elif mat.ndim >= 3:
         mat = mat.mean(axis=0)
 
-    plt.figure(figsize=(8, 4), dpi=120)
+    plt.figure(figsize=(7, 3.3), dpi=110)
     plt.imshow(mat, aspect="auto", cmap="RdBu_r", origin="upper")
-    plt.title(title)
+    plt.title(title, fontsize=10)
     plt.colorbar(shrink=0.7)
     plt.tight_layout()
 
     buf = io.BytesIO()
-    plt.savefig(buf, format="png")
-    plt.close()
-    buf.seek(0)
+    plt.savefig(buf, format="png", bbox_inches="tight")
+    plt.close(); buf.seek(0)
     return np.array(Image.open(buf))
 
 
-def encode_prompt_sd_xl(pipe, prompt: str, negative: str) -> Dict[str, torch.Tensor]:
-    """Return CLIP-L, CLIP-G (and negative) embeddings from SDXL pipeline."""
-    tok_l   = pipe.tokenizer(prompt,  max_length=77, padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device)
-    tok_g   = pipe.tokenizer_2(prompt,max_length=77, padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device)
-    ntok_l  = pipe.tokenizer(negative,  max_length=77, padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device)
-    ntok_g  = pipe.tokenizer_2(negative,max_length=77, padding="max_length", truncation=True, return_tensors="pt").input_ids.to(device)
+def encode_prompt_xl(pipe, prompt: str, negative: str) -> Dict[str, torch.Tensor]:
+    tok_l  = pipe.tokenizer  (prompt,  max_length=77, truncation=True,
+                              padding="max_length", return_tensors="pt").input_ids.to(device)
+    tok_g  = pipe.tokenizer_2(prompt,  max_length=77, truncation=True,
+                              padding="max_length", return_tensors="pt").input_ids.to(device)
+    ntok_l = pipe.tokenizer  (negative,max_length=77, truncation=True,
+                              padding="max_length", return_tensors="pt").input_ids.to(device)
+    ntok_g = pipe.tokenizer_2(negative,max_length=77, truncation=True,
+                              padding="max_length", return_tensors="pt").input_ids.to(device)
 
     with torch.no_grad():
-        clip_l = pipe.text_encoder(tok_l)[0]      # (1,77,768)
-        nclip_l= pipe.text_encoder(ntok_l)[0]
-        out_g  = pipe.text_encoder_2(tok_g, output_hidden_states=False)
-        clip_g, pooled = out_g[1], out_g[0]
-        nout_g = pipe.text_encoder_2(ntok_g, output_hidden_states=False)
-        nclip_g, npooled = nout_g[1], nout_g[0]
+        clip_l      = pipe.text_encoder(tok_l)[0]
+        neg_clip_l  = pipe.text_encoder(ntok_l)[0]
+
+        g_out       = pipe.text_encoder_2(tok_g,  output_hidden_states=False)
+        clip_g, pl  = g_out[1], g_out[0]
+        ng_out      = pipe.text_encoder_2(ntok_g, output_hidden_states=False)
+        neg_clip_g, npl = ng_out[1], ng_out[0]
 
     return {"clip_l": clip_l, "clip_g": clip_g,
-            "neg_l": nclip_l, "neg_g": nclip_g,
-            "pooled": pooled, "neg_pooled": npooled}
+            "neg_l":  neg_clip_l, "neg_g":  neg_clip_g,
+            "pooled": pl, "neg_pooled": npl}
 
 
-def adapter_forward(adapter, t5_seq, clip_seq, cfg):
-    with torch.no_grad():
-        out = adapter(t5_seq.float(), clip_seq.float())
-    # unify outputs
-    anchor, delta, log_sigma, *_, tau, g_pred, gate = (
-        out + (None,) * 8)[:8]   # pad to length 8
-    delta = delta * cfg["delta_scale"]
-    gate  = torch.sigmoid(gate * g_pred * cfg["gpred_scale"]) * cfg["gate_prob"]
-    final_delta = delta * cfg["strength"] * gate
-    mod = clip_seq + final_delta.to(dtype)
-
-    if cfg["sigma_scale"] > 0:
-        sigma = torch.exp(log_sigma * cfg["sigma_scale"])
-        mod += torch.randn_like(mod) * sigma.to(dtype)
-    if cfg["use_anchor"]:
-        mod = mod * (1 - gate) + anchor.to(dtype) * gate
-    if cfg["noise"] > 0:
-        mod += torch.randn_like(mod) * cfg["noise"]
-    return mod, final_delta, gate, g_pred, tau
-
-
-# ───────────────────────────────────────────────────────────────────────────
-# MAIN INFERENCE
-# ───────────────────────────────────────────────────────────────────────────
-def infer(prompt, negative_prompt,
-          adapter_l_file, adapter_g_file,
-          strength, delta_scale, sigma_scale,
-          gpred_scale, noise, gate_prob, use_anchor,
-          steps, cfg_scale, scheduler_name,
-          width, height, seed):
+# ─── INFERENCE ───────────────────────────────────────────────────────────
+def infer(prompt: str, negative_prompt: str,
+          adapter_l_file: str, adapter_g_file: str,
+          strength: float, delta_scale: float, sigma_scale: float,
+          gpred_scale: float, noise: float, gate_prob: float, use_anchor: bool,
+          steps: int, cfg_scale: float, scheduler_name: str,
+          width: int, height: int, seed: int):
 
     torch.cuda.empty_cache()
     _init_t5(); _init_pipe()
 
-    # scheduler
     if scheduler_name in SCHEDULERS:
         _pipe.scheduler = SCHEDULERS[scheduler_name].from_config(_pipe.scheduler.config)
 
-    # RNG
-    generator = None
-    if seed != -1:
-        generator = torch.Generator(device=device).manual_seed(seed)
-        torch.manual_seed(seed); np.random.seed(seed)
-
-    # T5 embeddings (semantic guidance)
-    t5_ids = _t5_tok(prompt, max_length=77, truncation=True, padding="max_length", return_tensors="pt").input_ids.to(device)
-    t5_seq = _t5_mod(t5_ids).last_hidden_state     # (1,77,768)
-
-    # CLIP embeddings from SDXL
-    embeds = encode_prompt_sd_xl(_pipe, prompt, negative_prompt)
-
-    # ------------------------------------------------------------------
-    # LOAD adapters (if any)
-    cfg_common = dict(
-        strength=strength, delta_scale=delta_scale, sigma_scale=sigma_scale,
-        gpred_scale=gpred_scale, noise=noise, gate_prob=gate_prob,
-        use_anchor=use_anchor,
-    )
+    generator = (torch.Generator(device=device).manual_seed(seed)
+                 if seed != -1 else None)
 
-    # --- STEP 0: build shift config -----------------------------------------
+    # build ShiftConfig (one per request)
     cfg_shift = ShiftConfig(
         prompt           = prompt,
         seed             = seed,
@@ -200,66 +146,76 @@ def infer(prompt, negative_prompt,
         use_anchor       = use_anchor,
         guidance_scale   = gpred_scale,
     )
-    
-    # --- STEP 1: encoder embeddings -----------------------------------------
+
+    # encoder (T5) embeddings
     t5_seq = ConditioningShifter.extract_encoder_embeddings(
         {"tokenizer": _t5_tok, "model": _t5_mod, "config": {"config": {}}},
         device, cfg_shift
     )
-    
-    # --- STEP 2: run adapters -----------------------------------------------
-    outputs = []
+
+    # CLIP embeddings
+    embeds = encode_prompt_xl(_pipe, prompt, negative_prompt)
+
+    # run adapters --------------------------------------------------------
+    outputs: List[AdapterOutput] = []
     if adapter_l_file and adapter_l_file != "None":
         ada_l = load_adapter(repo_l, adapter_l_file, conf_l, device)
         outputs.append(ConditioningShifter.run_adapter(
             ada_l, t5_seq, embeds["clip_l"],
             cfg_shift.guidance_scale, "clip_l", (0, 768)))
-    
+
     if adapter_g_file and adapter_g_file != "None":
         ada_g = load_adapter(repo_g, adapter_g_file, conf_g, device)
         outputs.append(ConditioningShifter.run_adapter(
             ada_g, t5_seq, embeds["clip_g"],
             cfg_shift.guidance_scale, "clip_g", (768, 2048)))
-    
-    # --- STEP 3: apply mods --------------------------------------------------
+
+    # apply modifications -------------------------------------------------
     clip_l_mod, clip_g_mod = embeds["clip_l"], embeds["clip_g"]
+    delta_viz = {"clip_l": torch.zeros_like(clip_l_mod),
+                 "clip_g": torch.zeros_like(clip_g_mod)}
+    gate_viz  = {"clip_l": torch.zeros_like(clip_l_mod[..., :1]),
+                 "clip_g": torch.zeros_like(clip_g_mod[..., :1])}
+
     for out in outputs:
-        tgt = clip_l_mod if out.adapter_type == "clip_l" else clip_g_mod
-        mod = ConditioningShifter.apply_modifications(tgt, [out], cfg_shift)
+        target = clip_l_mod if out.adapter_type == "clip_l" else clip_g_mod
+        mod    = ConditioningShifter.apply_modifications(target, [out], cfg_shift)
         if out.adapter_type == "clip_l":
             clip_l_mod = mod
         else:
             clip_g_mod = mod
+        delta_viz[out.adapter_type] = out.delta.detach()
+        gate_viz [out.adapter_type] = out.gate.detach()
 
-
-    # concatenate for SDXL
+    # prepare for SDXL ----------------------------------------------------
     prompt_embeds = torch.cat([clip_l_mod, clip_g_mod], dim=-1)
     neg_embeds    = torch.cat([embeds["neg_l"], embeds["neg_g"]], dim=-1)
 
-    # SDXL generation
     image = _pipe(
-        prompt_embeds            = prompt_embeds,
-        negative_prompt_embeds   = neg_embeds,
-        pooled_prompt_embeds     = embeds["pooled"],
+        prompt_embeds              = prompt_embeds,
+        negative_prompt_embeds     = neg_embeds,
+        pooled_prompt_embeds       = embeds["pooled"],
         negative_pooled_prompt_embeds = embeds["neg_pooled"],
-        num_inference_steps=steps, guidance_scale=cfg_scale,
-        width=width, height=height, generator=generator
+        num_inference_steps = steps,
+        guidance_scale      = cfg_scale,
+        width = width, height = height, generator = generator
     ).images[0]
 
-    # viz
-    delta_l_img = plot_heat(delta_l.squeeze(), "Δ CLIP-L")
-    gate_l_img  = plot_heat(gate_l.squeeze().mean(-1, keepdims=True), "Gate L")
-    delta_g_img = plot_heat(delta_g.squeeze(), "Δ CLIP-G")
-    gate_g_img  = plot_heat(gate_g.squeeze().mean(-1, keepdims=True), "Gate G")
+    # diagnostics ---------------------------------------------------------
+    delta_l_img = plot_heat(delta_viz["clip_l"].squeeze(), "Δ CLIP-L")
+    gate_l_img  = plot_heat(gate_viz ["clip_l"].squeeze().mean(-1, keepdims=True), "Gate L")
+    delta_g_img = plot_heat(delta_viz["clip_g"].squeeze(), "Δ CLIP-G")
+    gate_g_img  = plot_heat(gate_viz ["clip_g"].squeeze().mean(-1, keepdims=True), "Gate G")
+
+    stats_l = (f"τ̄_L = {outputs[0].tau.mean().item():.3f}"
+               if outputs and outputs[0].adapter_type == "clip_l" else "-")
+    stats_g = (f"τ̄_G = {outputs[-1].tau.mean().item():.3f}"
+               if len(outputs) > 1 and outputs[-1].adapter_type == "clip_g" else "-")
 
-    stats_l = f"g_pred_L={g_pred_l.item():.3f} | τ_L={tau_l.item():.3f}"
-    stats_g = f"g_pred_G={g_pred_g.item():.3f} | τ_G={tau_g.item():.3f}"
     return image, delta_l_img, gate_l_img, delta_g_img, gate_g_img, stats_l, stats_g
 
 
-# ───────────────────────────────────────────────────────────────────────────
-# GRADIO UI
-# ───────────────────────────────────────────────────────────────────────────
+# ─── GRADIO UI ────────────────────────────────────────────────────────────
 def create_interface():
     with gr.Blocks(title="SDXL Dual-Shunt Tester", theme=gr.themes.Soft()) as demo:
         gr.Markdown("# 🧠 SDXL Dual-Shunt Tester")
@@ -267,65 +223,65 @@ def create_interface():
         with gr.Row():
             with gr.Column(scale=1):
                 gr.Markdown("### Prompts")
-                prompt          = gr.Textbox(label="Prompt", lines=3,
-                                             value="a futuristic control station with holographic displays")
-                negative_prompt = gr.Textbox(label="Negative", lines=2,
-                                             value="blurry, low quality, distorted")
+                prompt   = gr.Textbox(label="Prompt", lines=3,
+                                      value="a futuristic control station with holographic displays")
+                negative = gr.Textbox(label="Negative", lines=2,
+                                      value="blurry, low quality, distorted")
 
                 gr.Markdown("### Adapters")
-                adapter_l = gr.Dropdown(["None"]+clip_l_opts, value="t5-vit-l-14-dual_shunt_caption.safetensors",
-                                         label="CLIP-L Adapter")
-                adapter_g = gr.Dropdown(["None"]+clip_g_opts, value="dual_shunt_omega_no_caption_noised_e1_step_10000.safetensors",
-                                         label="CLIP-G Adapter")
+                adapter_l = gr.Dropdown(["None"] + clip_l_opts,
+                                        value="t5-vit-l-14-dual_shunt_caption.safetensors",
+                                        label="CLIP-L Adapter")
+                adapter_g = gr.Dropdown(["None"] + clip_g_opts,
+                                        value="dual_shunt_omega_no_caption_noised_e1_step_10000.safetensors",
+                                        label="CLIP-G Adapter")
 
                 gr.Markdown("### Adapter Controls")
-                strength     = gr.Slider(0, 10, 4.0, 0.01, label="Strength")
-                delta_scale  = gr.Slider(-15, 15, 0.2, 0.1, label="Δ scale")
-                sigma_scale  = gr.Slider(0, 15, 0.1, 0.1, label="σ scale")
-                gpred_scale  = gr.Slider(0, 20, 2.0, 0.01, label="g_pred scale")
-                noise        = gr.Slider(0, 1, 0.55, 0.01, label="Extra noise")
-                gate_prob    = gr.Slider(0, 1, 0.27, 0.01, label="Gate prob")
-                use_anchor   = gr.Checkbox(True, label="Use anchor mix")
+                strength    = gr.Slider(0, 10, 4.0, 0.05, label="Strength")
+                delta_scale = gr.Slider(-15, 15, 0.2, 0.1, label="Δ scale")
+                sigma_scale = gr.Slider(0, 15, 0.1, 0.1, label="σ scale")
+                gpred_scale = gr.Slider(0, 20, 2.0, 0.05, label="Guidance scale")
+                noise       = gr.Slider(0, 1, 0.55, 0.01, label="Extra noise")
+                gate_prob   = gr.Slider(0, 1, 0.27, 0.01, label="Gate prob")
+                use_anchor  = gr.Checkbox(True, label="Use anchor mix")
 
                 gr.Markdown("### Generation")
                 with gr.Row():
-                    steps      = gr.Slider(1, 50, 20, 1, label="Steps")
-                    cfg_scale  = gr.Slider(1, 15, 7.5, 0.1, label="CFG")
-                scheduler   = gr.Dropdown(list(SCHEDULERS.keys()), value="DPM++ 2M", label="Scheduler")
+                    steps     = gr.Slider(1, 50, 20, 1, label="Steps")
+                    cfg_scale = gr.Slider(1, 15, 7.5, 0.1, label="CFG")
+                scheduler = gr.Dropdown(list(SCHEDULERS.keys()),
+                                        value="DPM++ 2M", label="Scheduler")
                 with gr.Row():
-                    width      = gr.Slider(512, 1536, 1024, 64, label="Width")
-                    height     = gr.Slider(512, 1536, 1024, 64, label="Height")
-                seed         = gr.Number(-1, label="Seed (-1=random)")
+                    width  = gr.Slider(512, 1536, 1024, 64, label="Width")
+                    height = gr.Slider(512, 1536, 1024, 64, label="Height")
+                seed = gr.Number(-1, label="Seed (-1 → random)", precision=0)
 
-                go_btn = gr.Button("🚀 Generate", variant="primary")
+                run_btn = gr.Button("🚀 Generate", variant="primary")
 
             with gr.Column(scale=1):
-                out_img   = gr.Image(label="Result", height=400)
-                gr.Markdown("### Adapter Diagnostics")
-                delta_l_i = gr.Image(label="Δ L", height=180)
-                gate_l_i  = gr.Image(label="Gate L", height=180)
-                delta_g_i = gr.Image(label="Δ G", height=180)
-                gate_g_i  = gr.Image(label="Gate G", height=180)
-                stats_l   = gr.Textbox(label="Stats L", interactive=False)
-                stats_g   = gr.Textbox(label="Stats G", interactive=False)
+                out_img  = gr.Image(label="Result", height=400)
+                gr.Markdown("### Diagnostics")
+                delta_l  = gr.Image(label="Δ L", height=180)
+                gate_l   = gr.Image(label="Gate L", height=180)
+                delta_g  = gr.Image(label="Δ G", height=180)
+                gate_g   = gr.Image(label="Gate G", height=180)
+                stats_l  = gr.Textbox(label="Stats L", interactive=False)
+                stats_g  = gr.Textbox(label="Stats G", interactive=False)
 
         def _run(*args):
-            pl   , npl  = args[0], args[1]
-            al, ag      = (None if v=="None" else v for v in args[2:4])
+            pl, npl = args[0], args[1]
+            al, ag  = (None if v == "None" else v for v in args[2:4])
             return infer(pl, npl, al, ag, *args[4:])
 
-        go_btn.click(
-            _run,
-            inputs=[prompt, negative_prompt, adapter_l, adapter_g,
-                    strength, delta_scale, sigma_scale, gpred_scale,
-                    noise, gate_prob, use_anchor, steps, cfg_scale,
-                    scheduler, width, height, seed],
-            outputs=[out_img, delta_l_i, gate_l_i, delta_g_i, gate_g_i,
-                     stats_l, stats_g]
+        run_btn.click(
+            fn=_run,
+            inputs=[prompt, negative, adapter_l, adapter_g, strength, delta_scale,
+                    sigma_scale, gpred_scale, noise, gate_prob, use_anchor, steps,
+                    cfg_scale, scheduler, width, height, seed],
+            outputs=[out_img, delta_l, gate_l, delta_g, gate_g, stats_l, stats_g]
         )
     return demo
 
 
-# ───────────────────────────────────────────────────────────────────────────
 if __name__ == "__main__":
     create_interface().launch()