Update app.py

app.py

@@ -3,6 +3,7 @@ from gradio_client import Client
 import os
 from dotenv import load_dotenv
 import warnings
+import torch
 
 # Suppress warnings
 warnings.filterwarnings('ignore', category=UserWarning)
@@ -11,12 +12,20 @@ warnings.filterwarnings('ignore', category=UserWarning)
 load_dotenv()
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# Force CPU mode
+os.environ['CUDA_VISIBLE_DEVICES'] = ''
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+else:
+    device = torch.device('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,
+            cpu_only=True  # Force CPU mode
         )
         result = client.predict(
             prompt,
@@ -313,12 +322,14 @@ demo = gr.Interface(
     cache_examples=True
 )
 
-# Launch app
 if __name__ == "__main__":
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True,
-        debug=True
-    )
+    try:
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            show_error=True,
+            debug=True
+        )
+    except Exception as e:
+        print(f"Error launching app: {str(e)}")