Update models.py

models.py

@@ -5,70 +5,50 @@ from diffusers import StableDiffusionPipeline, DiffusionPipeline
 from huggingface_hub import login
 from typing import Tuple
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-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)
+def load_models() -> Tuple[StableDiffusionPipeline, DiffusionPipeline, None]:
     try:
-        # Authentication setup
+        # Device and precision configuration
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        dtype = torch.float16 if device.type == "cuda" else torch.float32
+        
+        # Authentication
         hf_token = os.getenv("HF_TOKEN")
         if hf_token:
             login(token=hf_token)
-            logger.info("HF authentication successful")
-        else:
-            logger.warning("Proceeding without HF authentication")
-
-        # 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 optimizations
-        logger.info("Loading text-to-image model...")
+        # Text-to-image model
+        logger.info(f"Loading text-to-image model on {device} with {dtype}")
         text_to_image = StableDiffusionPipeline.from_pretrained(
             "runwayml/stable-diffusion-v1-5",
-            torch_dtype=torch.float16,
+            torch_dtype=dtype,
             use_safetensors=True,
-            safety_checker=None,
-            variant="fp16",
-            use_auth_token=hf_token if hf_token else None
+            safety_checker=None
         )
+        text_to_image = text_to_image.to(device)
         
-        # Enable memory optimizations
-        if torch.cuda.is_available():
-            text_to_image = text_to_image.to("cuda")
+        if device.type == "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 image-to-video model
-        logger.info("Loading image-to-video model...")
+        # Image-to-video model
+        logger.info(f"Loading video model on {device} with {dtype}")
         image_to_video = DiffusionPipeline.from_pretrained(
             "cerspense/zeroscope_v2_576w",
-            torch_dtype=torch.float16,
-            use_auth_token=hf_token if hf_token else None
+            torch_dtype=dtype
         )
-
-        # Video model optimizations
-        if torch.cuda.is_available():
-            image_to_video = image_to_video.to("cuda")
+        image_to_video = image_to_video.to(device)
+        
+        if device.type == "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  # TTS placeholder
+        return text_to_image, image_to_video, None
 
     except Exception as e:
-        logger.error(f"Model loading failed: {str(e)}")
-        raise RuntimeError("Model initialization error - check logs") from e
+        logger.error(f"Model load failed: {str(e)}")
+        raise