Update app.py

app.py

@@ -24,7 +24,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 pipe = DiffusionPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell",
     torch_dtype=dtype,
-    device_map="auto",
+    device_map="balanced" if torch.cuda.is_available() else None,
     use_safetensors=True
 ).to(device)
 
@@ -70,6 +70,7 @@ def generate_diagram(prompt, seed=42, randomize_seed=False, width=1024, height=1
 
 
 
+
 # Enhanced examples with more detailed prompts and specific styling
 EXAMPLES = [
     {