debugging app.py

app.py

@@ -40,6 +40,23 @@ def process_image_and_text(image, scale, seed, text):
 
     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)
+        
+    missing_keys, unexpected_keys = pipe.transformer.load_state_dict(key_changed_blended_attn_weights, strict=False)
+
     out = pipe(
         prompt=text,
         height=512,