Update models.py

models.py

@@ -9,8 +9,11 @@ from typing import Tuple
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def load_models() -> Tuple[StableDiffusionPipeline, DiffusionPipeline, None]:
-    """Load and configure AI models with memory optimizations."""
+def load_models():
+    # Existing model loading logic
+    # Add this safety check
+    if not hasattr(torch, 'cuda') or not torch.cuda.is_available():
+        torch.set_flush_denormal(True)
     try:
         # Authentication setup
         hf_token = os.getenv("HF_TOKEN")