final_app

app.py

@@ -34,6 +34,11 @@ def timeout_handler(signum, frame):
 # Load the model
 def load_model():
     try:
+        # Initialize signal handler only on Unix-like systems
+        if TORCH_DEVICE == "cpu" and hasattr(signal, 'SIGALRM'):
+            signal.signal(signal.SIGALRM, timeout_handler)
+            signal.alarm(2100)  # 15 minutes timeout for model loading
+        
         pipe = DiffusionPipeline.from_pretrained(
             "CompVis/stable-diffusion-v1-4",
             torch_dtype=torch.float16 if TORCH_DEVICE == "cuda" else torch.float32,
@@ -63,9 +68,17 @@ def load_model():
                     print(f"Warning: Could not load textual inversion concept {concept}: {e}")
         except Exception as e:
             print(f"Warning: Could not load textual inversion concepts: {e}")
+        
+        # Clear the alarm if set
+        if TORCH_DEVICE == "cpu" and hasattr(signal, 'SIGALRM'):
+            signal.alarm(0)
             
         return pipe
     except Exception as e:
+        # Clear the alarm if set
+        if TORCH_DEVICE == "cpu" and hasattr(signal, 'SIGALRM'):
+            signal.alarm(0)
+        
         print(f"Error loading model: {e}")
         traceback.print_exc()
         raise