Update app.py

app.py

@@ -14,7 +14,7 @@ from configs import T5_SHUNT_REPOS
 # ─── Device & Model Setup ─────────────────────────────────────
 # Don't initialize CUDA here for ZeroGPU compatibility
 device = None  # Will be set inside the GPU function
-dtype = torch.float32
+dtype = torch.float16
 
 # Don't load models here - will load inside GPU function
 t5_tok = None
@@ -66,7 +66,7 @@ def plot_heat(mat, title):
     return buf
 
 # ─── SDXL Text Encoding ───────────────────────────────────────
-def encode_sdxl_prompt(pipe, prompt, negative_prompt="", device=device):
+def encode_sdxl_prompt(prompt, negative_prompt=""):
     """Generate proper CLIP-L and CLIP-G embeddings using SDXL's text encoders"""
     
     # Tokenize for both encoders
@@ -128,6 +128,10 @@ def encode_sdxl_prompt(pipe, prompt, negative_prompt="", device=device):
         "neg_pooled": neg_pooled_embeds
     }
 
+# ─── Inference ────────────────────────────────────────────────
+@torch.no_grad()
+def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noise, gate_prob, 
+          use_anchor, steps, cfg_scale, scheduler_name, width, height, seed):
     
 # ─── Inference ────────────────────────────────────────────
 @spaces.GPU
@@ -186,8 +190,7 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
     
     # Apply CLIP-L adapter
     if adapter_l is not None:
-        clip_l_in = clip_embeds["clip_l"].to(torch.float32)
-        anchor_l, delta_l, log_sigma_l, attn_l1, attn_l2, tau_l, g_pred_l, gate_l = adapter_l(t5_seq, clip_l_in)
+        anchor_l, delta_l, log_sigma_l, attn_l1, attn_l2, tau_l, g_pred_l, gate_l = adapter_l(t5_seq, clip_embeds["clip_l"])
         gate_l_scaled = gate_l * gate_prob
         delta_l_final = delta_l * strength * gate_l_scaled
         clip_l_mod = clip_embeds["clip_l"] + delta_l_final
@@ -204,10 +207,7 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
     
     # Apply CLIP-G adapter  
     if adapter_g is not None:
-        # Float32 adapter input
-        clip_g_in = clip_embeds["clip_g"].to(torch.float32)
-
-        anchor_g, delta_g, log_sigma_g, attn_g1, attn_g2, tau_g, g_pred_g, gate_g = adapter_g(t5_seq, clip_g_in)
+        anchor_g, delta_g, log_sigma_g, attn_g1, attn_g2, tau_g, g_pred_g, gate_g = adapter_g(t5_seq, clip_embeds["clip_g"])
         gate_g_scaled = gate_g * gate_prob
         delta_g_final = delta_g * strength * gate_g_scaled
         clip_g_mod = clip_embeds["clip_g"] + delta_g_final
@@ -243,9 +243,9 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
     return (
         image,
         plot_heat(delta_l_final.squeeze().cpu().numpy(), "Δ CLIP-L"),
-        plot_heat(gate_l_scaled.squeeze().cpu().numpy(), "Gate CLIP-L"), 
+        plot_heat(gate_l_scaled.squeeze().cpu().numpy().mean(axis=-1), "Gate CLIP-L"), 
         plot_heat(delta_g_final.squeeze().cpu().numpy(), "Δ CLIP-G"),
-        plot_heat(gate_g_scaled.squeeze().cpu().numpy(), "Gate CLIP-G"),
+        plot_heat(gate_g_scaled.squeeze().cpu().numpy().mean(axis=-1), "Gate CLIP-G"),
         f"g_pred_l: {g_pred_l.mean().item():.3f}, τ_l: {tau_l.mean().item():.3f}",
         f"g_pred_g: {g_pred_g.mean().item():.3f}, τ_g: {tau_g.mean().item():.3f}"
     )