add torchvision to requirement

app.py

@@ -12,10 +12,49 @@ from src.utils_sample import set_seed, resize_and_add_margin
 import os
 
 
+dtype = torch.bfloat16
+token = os.environ.get("HF_TOKEN")
+
+pipe = None
+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, pipe):
+def process_image_and_text(image, scale, seed, text):
     set_seed(seed)
+    print("execution_device 1", pipe._execution_device)
+    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 = pipe.to(dtype)
+    pipe = pipe.to("cuda")
+    print("execution_device 2", pipe._execution_device)
+
+    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)
+
     # image = Image.open(img_path).convert('RGB') 
     image = resize_and_add_margin(image, target_size=512)
     
@@ -108,16 +147,6 @@ header = """
 
 def create_app():
     
-    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")
-
     with gr.Blocks() as app:
         gr.Markdown(header, elem_id="header")
         with gr.Row(equal_height=False):
@@ -147,39 +176,9 @@ def create_app():
                 label="Examples",
             )
         
-        print("execution_device 1", pipe._execution_device)
-        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 = pipe.to(dtype)
-        pipe = pipe.to("cuda")
-        print("execution_device 2", pipe._execution_device)
-
-        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)
-
-
         submit_btn.click(
             fn=process_image_and_text,
-            inputs=[original_image, scale, seed, text, pipe],
+            inputs=[original_image, scale, seed, text],
             outputs=output_image,
         )
         

requirements.txt

@@ -1,9 +1,8 @@
-torch
-torchvision
 transformers
 protobuf
 sentencepiece
 accelerate
 einops
 huggingface_hub
+torchvision
 git+https://github.com/WonwoongCho/diffusers@main#egg=diffusers