debugging app.py

app.py

@@ -21,14 +21,25 @@ pipe = FluxPipeline.from_pretrained(
 )
 pipe = pipe.to("cuda")
 
+
 @spaces.GPU
-def process_image_and_text(image, text, seed):
+def process_image_and_text(image, text, seed, scale):
     set_seed(seed)
-    # image = Image.open(img_path).convert('RGB') 
     image = resize_and_add_margin(image, target_size=512)
-    
     image_list = [image]
-    
+
+    # Dynamically set attention processors using user-specified scale
+    blended_attn_procs = {}
+    for name, _ in pipe.transformer.attn_processors.items():
+        if "single" in name:
+            processor = FluxBlendedAttnProcessor2_0(3072, ba_scale=float(scale), num_ref=1)
+            processor = processor.to(device="cuda", dtype=dtype)
+            blended_attn_procs[name] = processor
+        else:
+            blended_attn_procs[name] = pipe.transformer.attn_processors[name]
+
+    pipe.transformer.set_attn_processor(blended_attn_procs)
+
     out = pipe(
         prompt=text,
         height=512,
@@ -139,32 +150,6 @@ def create_app():
                 output_image = gr.Image(type="pil", elem_id="output")
 
 
-        blended_attn_procs = {}
-        for name, _ in pipe.transformer.attn_processors.items():
-            if "single" in name:
-                blended_attn_procs[name] = FluxBlendedAttnProcessor2_0(3072, ba_scale=scale, num_ref=1).to(device="cuda", dtype=dtype)
-            else:
-                blended_attn_procs[name] = pipe.transformer.attn_processors[name]
-
-        pipe.transformer.set_attn_processor(blended_attn_procs)
-
-        model_path = hf_hub_download(
-            repo_id="WonwoongCho/IT-Blender",
-            filename="FLUX/it-blender.bin",
-            token=token
-        )
-        
-        pretrained_blended_attn_weights = torch.load(model_path, map_location=pipe._execution_device)
-
-        key_changed_blended_attn_weights = {}
-        for key, value in pretrained_blended_attn_weights.items():
-            block_idx = int(key.split(".")[0]) - 21
-            k_or_v = key.split("_")[2]
-            changed_key = f'single_transformer_blocks.{block_idx}.attn.processor.blended_attention_{k_or_v}_proj.weight'
-            key_changed_blended_attn_weights[changed_key] = value.to(dtype=dtype, device="cuda")
-            
-        missing_keys, unexpected_keys = pipe.transformer.load_state_dict(key_changed_blended_attn_weights, strict=False)
-
         with gr.Row():
             examples = gr.Examples(
                 examples=get_samples(),
@@ -174,7 +159,7 @@ def create_app():
 
         submit_btn.click(
             fn=process_image_and_text,
-            inputs=[original_image, text, seed],
+            inputs=[original_image, text, seed, scale],
             outputs=output_image,
         )