debugging app.py

app.py

@@ -11,6 +11,15 @@ from src.attention_processor import FluxBlendedAttnProcessor2_0
 from src.utils_sample import set_seed, resize_and_add_margin
 import os
 
+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("cuda")
 
 @spaces.GPU
 def process_image_and_text(image, text, seed, pipeline):
@@ -105,18 +114,7 @@ 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)
+def create_app():
 
     with gr.Blocks() as app:
         gr.Markdown(header, elem_id="header")
@@ -167,7 +165,7 @@ def create_app(
             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)
+        pipe = pipe.to("cuda")
         
 
         with gr.Row():
@@ -188,5 +186,5 @@ def create_app(
 
 
 if __name__ == "__main__":
-    demo = create_app(device="cuda")
+    demo = create_app()
     demo.launch(debug=True, ssr_mode=False)