debugging app.py

app.py

@@ -11,47 +11,9 @@ from src.attention_processor import FluxBlendedAttnProcessor2_0
 from src.utils_sample import set_seed, resize_and_add_margin
 import os
 
-pipe = None
-dtype = torch.bfloat16
-token = os.environ.get("HF_TOKEN")
-
-pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev", 
-    torch_dtype=dtype, 
-    use_auth_token=token
-)
-pipe = pipe.to("cuda")
 
 @spaces.GPU
-def process_image_and_text(image, scale, seed, text):
-    set_seed(seed)
-    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)
-        else:
-            blended_attn_procs[name] = pipe.transformer.attn_processors[name]
-
-    pipe.transformer.set_attn_processor(blended_attn_procs)
-    pipe.to(dtype)
-
-    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)
-        
-    missing_keys, unexpected_keys = pipe.transformer.load_state_dict(key_changed_blended_attn_weights, strict=False)
-    pipe = pipe.to("cuda")
+def process_image_and_text(image, text, pipe):
     
     # image = Image.open(img_path).convert('RGB') 
     image = resize_and_add_margin(image, target_size=512)
@@ -146,7 +108,16 @@ header = """
 def create_app(
     device: str = "cuda",
 ):
-    
+    dtype = torch.bfloat16
+    token = os.environ.get("HF_TOKEN")
+
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-dev", 
+        torch_dtype=dtype, 
+        token=token
+    )
+    pipe = pipe.to(device)
+
     with gr.Blocks() as app:
         gr.Markdown(header, elem_id="header")
         with gr.Row(equal_height=False):
@@ -169,16 +140,47 @@ def create_app(
             with gr.Column(variant="panel", elem_classes="outputPanel"):
                 output_image = gr.Image(type="pil", elem_id="output")
 
+
+        set_seed(seed)
+        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)
+            else:
+                blended_attn_procs[name] = pipe.transformer.attn_processors[name]
+
+        pipe.transformer.set_attn_processor(blended_attn_procs)
+        pipe.to(dtype)
+
+        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)
+            
+        missing_keys, unexpected_keys = pipe.transformer.load_state_dict(key_changed_blended_attn_weights, strict=False)
+        pipe = pipe.to(device)
+        
+
         with gr.Row():
             examples = gr.Examples(
                 examples=get_samples(),
-                inputs=[original_image, scale, seed, text],
+                inputs=[original_image, text],
                 label="Examples",
             )
         
         submit_btn.click(
             fn=process_image_and_text,
-            inputs=[original_image, scale, seed, text],
+            inputs=[original_image, text, pipe],
             outputs=output_image,
         )