Update app.py

app.py

@@ -17,7 +17,12 @@ base_path = "black-forest-labs/FLUX.1-dev"
 lora_base_path = "./models"
 
 
-pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
+pipe = FluxPipeline.from_pretrained(
+    base_path, 
+    torch_dtype=torch.bfloat16, 
+    device_map="balanced", 
+    max_memory={0:"16GB", 1:"16GB"}
+)
 transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16)
 pipe.transformer = transformer
 pipe.to("cuda")