Update app.py

app.py

@@ -2,11 +2,36 @@ import gradio as gr
 from gradio_client import Client
 import os
 from dotenv import load_dotenv
+import torch
 
 # Load environment variables
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# Check CUDA availability and set device
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+def generate_diagram(prompt, width=1024, height=1024):
+    """Generate a diagram using FLUX AI"""
+    try:
+        client = Client(
+            "black-forest-labs/FLUX.1-schnell",
+            hf_token=HF_TOKEN,
+        )
+        result = client.predict(
+            prompt,
+            1872187377,  # seed
+            False,       # randomize_seed
+            width,
+            height,
+            4,          # num_inference_steps
+            api_name="/infer"
+        )
+        return result
+    except Exception as e:
+        raise gr.Error(f"Error generating diagram: {str(e)}")
+
+
 # Enhanced examples with more detailed prompts and specific styling
 EXAMPLES = [
     {
@@ -202,22 +227,7 @@ EXAMPLES = [
     }
 ]
 
-def generate_diagram(prompt, width=1024, height=1024):
-    """Generate a diagram using FLUX AI"""
-    try:
-        client = Client("black-forest-labs/FLUX.1-schnell", hf_token=HF_TOKEN)
-        result = client.predict(
-            prompt,
-            1872187377,  # seed
-            False,       # randomize_seed
-            width,
-            height,
-            4,          # num_inference_steps
-            api_name="/infer"
-        )
-        return result
-    except Exception as e:
-        raise gr.Error(f"Error generating diagram: {str(e)}")
+
 
 # Convert examples to Gradio format
 GRADIO_EXAMPLES = [
@@ -251,61 +261,39 @@ demo = gr.Interface(
     ],
     outputs=gr.Image(type="filepath", label="Generated Diagram"),
     title="🎨 FLUX Diagram Generator",
-    description="""Generate beautiful hand-drawn style diagrams using FLUX AI. 
-    Create intricate mind maps with custom themes and detailed iconography.""",
+    description="Generate beautiful hand-drawn style diagrams using FLUX AI",
     article="""
     ### Tips for Better Results
-    - Use descriptive style hints (e.g., "vibrant colors", "flowing connections")
-    - Include specific icon sizes (e.g., "~60px")
-    - Add theme-specific patterns (e.g., "circuit board patterns", "wave patterns")
-    - Use consistent hierarchy and formatting
-    - Combine both functional and decorative elements
+    - Use clear hierarchical structures
+    - Include icon descriptions in brackets
+    - Keep text concise and meaningful
+    - Use consistent formatting
     
     ### Template Structure
     ```
     MAIN TOPIC
-    ├── PRIMARY [Icon ~60px]
-    │   ├── SUB-A [Specific Icon]
-    │   └── SUB-B [Detailed Icon]
-    ├── SECONDARY [Icon ~55px]
-    │   ├── SUB-C [Themed Icon]
-    │   └── SUB-D [Custom Icon]
-    └── TERTIARY
-        ├── SUB-E [Dynamic Icon]
-        └── SUB-F [Unique Icon]
+    ├── SUBTOPIC 1 [Icon]
+    ├── SUBTOPIC 2 [Icon]
+    └── SUBTOPIC 3
+        ├── DETAIL 1 [Icon]
+        └── DETAIL 2 [Icon]
     ```
     """,
-    examples=GRADIO_EXAMPLES,
+    examples=EXAMPLES,
     cache_examples=True,
-    theme=gr.themes.Soft(),
-    css="""
-        .gradio-container {
-            font-family: 'Inter', sans-serif;
-        }
-        .gr-prose {
-            max-width: 850px;
-            margin: auto;
-        }
-        .gr-prose code {
-            background-color: #f3f4f6;
-            padding: 0.2em 0.4em;
-            border-radius: 0.3em;
-            font-size: 0.95em;
-        }
-        .gr-prose h3 {
-            margin-top: 1.5em;
-            color: #2563eb;
-        }
-    """
+    theme=gr.themes.Soft()
 )
 
-# Launch app
+# Launch app with error handling
 if __name__ == "__main__":
-    # Set HF space configs
-    demo.queue(max_size=15)  # Increased queue size
-    demo.launch(
-        share=False,
-        show_error=True,
-        server_name="0.0.0.0",
-        server_port=7860,
-    )
+    try:
+        demo.queue(max_size=10)
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            show_error=True,
+            enable_queue=True
+        )
+    except Exception as e:
+        print(f"Error launching app: {str(e)}")