Update models.py

models.py

@@ -1,51 +1,71 @@
 import torch
 import logging
+import os
 from diffusers import StableDiffusionPipeline, DiffusionPipeline
 from huggingface_hub import login
-import os
+from typing import Tuple
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-def load_models():
+def load_models() -> Tuple[StableDiffusionPipeline, DiffusionPipeline, None]:
+    """Load and configure AI models with memory optimizations."""
     try:
-        # Authentication
-        hf_token = os.getenv("HF_TOKEN", "")
+        # Authentication setup
+        hf_token = os.getenv("HF_TOKEN")
         if hf_token:
             login(token=hf_token)
+            logger.info("HF authentication successful")
         else:
-            logger.warning("No HF_TOKEN found, using anonymous access")
+            logger.warning("Proceeding without HF authentication")
 
-        # Memory optimization config
-        torch.backends.cudnn.benchmark = True
-        torch.set_flush_denormal(True)
+        # Configure Torch for optimal performance
+        torch.set_grad_enabled(False)
+        if torch.cuda.is_available():
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.benchmark = True
 
-        # Load text-to-image model with memory optimization
+        # Load text-to-image model with optimizations
         logger.info("Loading text-to-image model...")
         text_to_image = StableDiffusionPipeline.from_pretrained(
-            "runwayml/stable-diffusion-v1-5",  # Lighter variant
+            "runwayml/stable-diffusion-v1-5",
             torch_dtype=torch.float16,
             use_safetensors=True,
             safety_checker=None,
-            requires_safety_checker=False
+            variant="fp16",
+            use_auth_token=hf_token if hf_token else None
         )
-        text_to_image = text_to_image.to("cpu")
-        text_to_image.enable_attention_slicing()
-        text_to_image.enable_sequential_cpu_offload()
+        
+        # Enable memory optimizations
+        if torch.cuda.is_available():
+            text_to_image = text_to_image.to("cuda")
+            text_to_image.enable_xformers_memory_efficient_attention()
+            text_to_image.enable_model_cpu_offload()
+        else:
+            text_to_image = text_to_image.to("cpu")
+            text_to_image.enable_attention_slicing()
 
-        # Load video model with reduced parameters
+        # Load image-to-video model
         logger.info("Loading image-to-video model...")
         image_to_video = DiffusionPipeline.from_pretrained(
-            "cerspense/zeroscope_v2_576w",  # Smaller video model
-            torch_dtype=torch.float16
+            "cerspense/zeroscope_v2_576w",
+            torch_dtype=torch.float16,
+            use_auth_token=hf_token if hf_token else None
         )
-        image_to_video = image_to_video.to("cpu")
-        image_to_video.enable_attention_slicing(1)
+
+        # Video model optimizations
+        if torch.cuda.is_available():
+            image_to_video = image_to_video.to("cuda")
+            image_to_video.enable_xformers_memory_efficient_attention()
+            image_to_video.enable_model_cpu_offload()
+        else:
+            image_to_video = image_to_video.to("cpu")
+            image_to_video.enable_attention_slicing()
 
         logger.info("All models loaded successfully")
-        return text_to_image, image_to_video, None
+        return text_to_image, image_to_video, None  # TTS placeholder
 
     except Exception as e:
         logger.error(f"Model loading failed: {str(e)}")
-        raise RuntimeError("Could not initialize models") from e
+        raise RuntimeError("Model initialization error - check logs") from e