Update app.py

app.py

@@ -133,8 +133,9 @@ def encode_sdxl_prompt(pipe, prompt, negative_prompt, device):
 
 # ─── Main Inference Function ──────────────────────────────────
 @spaces.GPU
-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):
+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):
     
     global t5_tok, t5_mod, pipe
     device = torch.device("cuda")
@@ -193,10 +194,17 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
                 # Handle different return formats
                 anchor_l = adapter_output[0]
                 delta_l = adapter_output[1]
+                log_sigma_l = adapter_output[2] if len(adapter_output) > 2 else torch.zeros_like(delta_l)
                 gate_l = adapter_output[-1] if len(adapter_output) > 2 else torch.ones_like(delta_l)
                 tau_l = adapter_output[-2] if len(adapter_output) > 6 else torch.tensor(1.0)
                 g_pred_l = adapter_output[-3] if len(adapter_output) > 6 else torch.tensor(1.0)
             
+            # Scale delta values
+            delta_l = delta_l * delta_scale
+            
+            # Apply g_pred scaling to gate
+            gate_l = gate_l * g_pred_l * gpred_scale
+            
             # Apply gate scaling
             gate_l_scaled = torch.sigmoid(gate_l) * gate_prob
             
@@ -206,11 +214,16 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
             # Apply delta to embeddings
             clip_l_mod = clip_embeds["clip_l"] + delta_l_final.to(dtype)
             
+            # Apply sigma-based noise if specified
+            if sigma_scale > 0:
+                sigma_l = torch.exp(log_sigma_l * sigma_scale)
+                clip_l_mod += torch.randn_like(clip_l_mod) * sigma_l.to(dtype)
+            
             # Apply anchor mixing if enabled
             if use_anchor:
                 clip_l_mod = clip_l_mod * (1 - gate_l_scaled.to(dtype)) + anchor_l.to(dtype) * gate_l_scaled.to(dtype)
             
-            # Add noise if specified
+            # Add additional noise if specified
             if noise > 0:
                 clip_l_mod += torch.randn_like(clip_l_mod) * noise
     else:
@@ -233,10 +246,17 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
                 # Handle different return formats
                 anchor_g = adapter_output[0]
                 delta_g = adapter_output[1]
+                log_sigma_g = adapter_output[2] if len(adapter_output) > 2 else torch.zeros_like(delta_g)
                 gate_g = adapter_output[-1] if len(adapter_output) > 2 else torch.ones_like(delta_g)
                 tau_g = adapter_output[-2] if len(adapter_output) > 6 else torch.tensor(1.0)
                 g_pred_g = adapter_output[-3] if len(adapter_output) > 6 else torch.tensor(1.0)
             
+            # Scale delta values
+            delta_g = delta_g * delta_scale
+            
+            # Apply g_pred scaling to gate
+            gate_g = gate_g * g_pred_g * gpred_scale
+            
             # Apply gate scaling
             gate_g_scaled = torch.sigmoid(gate_g) * gate_prob
             
@@ -246,11 +266,16 @@ def infer(prompt, negative_prompt, adapter_l_file, adapter_g_file, strength, noi
             # Apply delta to embeddings
             clip_g_mod = clip_embeds["clip_g"] + delta_g_final.to(dtype)
             
+            # Apply sigma-based noise if specified
+            if sigma_scale > 0:
+                sigma_g = torch.exp(log_sigma_g * sigma_scale)
+                clip_g_mod += torch.randn_like(clip_g_mod) * sigma_g.to(dtype)
+            
             # Apply anchor mixing if enabled
             if use_anchor:
                 clip_g_mod = clip_g_mod * (1 - gate_g_scaled.to(dtype)) + anchor_g.to(dtype) * gate_g_scaled.to(dtype)
             
-            # Add noise if specified
+            # Add additional noise if specified
             if noise > 0:
                 clip_g_mod += torch.randn_like(clip_g_mod) * noise
     else:
@@ -330,8 +355,11 @@ def create_interface():
                 
                 # Controls
                 gr.Markdown("### 🎛️ Adapter Controls")
-                strength = gr.Slider(0.0, 5.0, value=1.0, step=0.1, label="Adapter Strength")
-                noise = gr.Slider(0.0, 1.0, value=0.0, step=0.05, label="Noise Injection")
+                strength = gr.Slider(0.0, 5.0, value=1.0, step=0.01, label="Adapter Strength")
+                delta_scale = gr.Slider(-15.0, 15.0, value=1.0, step=0.1, label="Delta Scale", info="Scales the delta values, recommended 1")
+                sigma_scale = gr.Slider(-15.0, 15.0, value=1.0, step=0.1, label="Sigma Scale", info="Scales the noise variance, recommended 1")
+                gpred_scale = gr.Slider(0.0, 10.0, value=2.0, step=0.1, label="G-Pred Scale", info="Scales the gate prediction, recommended 2")
+                noise = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="Noise Injection")
                 gate_prob = gr.Slider(0.0, 1.0, value=1.0, step=0.05, label="Gate Probability")
                 use_anchor = gr.Checkbox(label="Use Anchor Points", value=True)
                 
@@ -385,8 +413,9 @@ def create_interface():
         generate_btn.click(
             fn=run_generation,
             inputs=[
-                prompt, negative_prompt, adapter_l, adapter_g, strength, noise, gate_prob,
-                use_anchor, steps, cfg_scale, scheduler_name, width, height, seed
+                prompt, negative_prompt, adapter_l, adapter_g, strength, delta_scale,
+                sigma_scale, gpred_scale, noise, gate_prob, use_anchor, steps, cfg_scale,
+                scheduler_name, width, height, seed
             ],
             outputs=[output_image, delta_l_img, gate_l_img, delta_g_img, gate_g_img, stats_l_text, stats_g_text]
         )