add torchvision to requirement

app.py

@@ -11,29 +11,37 @@ 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, 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)
+    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 = {}
@@ -45,6 +53,8 @@ def process_image_and_text(image, scale, seed, text, pipe):
         
     missing_keys, unexpected_keys = pipe.transformer.load_state_dict(key_changed_blended_attn_weights, strict=False)
 
+    pipe = pipe.to("cuda")
+
     # image = Image.open(img_path).convert('RGB') 
     image = resize_and_add_margin(image, target_size=512)
     
@@ -137,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):
@@ -175,10 +175,10 @@ def create_app():
                 inputs=[original_image, scale, seed, text],
                 label="Examples",
             )
-
+        
         submit_btn.click(
             fn=process_image_and_text,
-            inputs=[original_image, scale, seed, text, pipe],
+            inputs=[original_image, scale, seed, text],
             outputs=output_image,
         )