Update models.py

models.py

@@ -4,6 +4,7 @@ UPDATED VERSION with proper InstantID pipeline support
 """
 import torch
 import time
+import os
 from diffusers import (
     ControlNetModel,
     AutoencoderKL,
@@ -208,6 +209,31 @@ def optimize_pipeline(pipe):
             print("  [OK] xformers enabled")
         except Exception as e:
             print(f"  [INFO] xformers not available: {e}")
+    
+    # Additional optimizations for memory efficiency
+    if hasattr(pipe, 'enable_vae_slicing'):
+        pipe.enable_vae_slicing()
+        print("  [OK] VAE slicing enabled for memory efficiency")
+    
+    if hasattr(pipe, 'enable_vae_tiling'):
+        pipe.enable_vae_tiling()
+        print("  [OK] VAE tiling enabled for memory efficiency")
+    
+    # For Zero GPU environments (HuggingFace Spaces)
+    if os.environ.get("SPACE_ID"):
+        print("  [INFO] Detected HuggingFace Spaces environment")
+        if hasattr(pipe, 'enable_model_cpu_offload'):
+            # Don't enable CPU offload if using Zero GPU
+            # as it conflicts with @spaces.GPU decorator
+            print("  [INFO] Zero GPU environment - skipping CPU offload (handled by @spaces.GPU)")
+    
+    # Ensure dtype consistency
+    if hasattr(pipe, 'vae') and pipe.vae is not None:
+        pipe.vae = pipe.vae.to(dtype=dtype)
+    if hasattr(pipe, 'text_encoder') and pipe.text_encoder is not None:
+        pipe.text_encoder = pipe.text_encoder.to(dtype=dtype)
+    if hasattr(pipe, 'text_encoder_2') and pipe.text_encoder_2 is not None:
+        pipe.text_encoder_2 = pipe.text_encoder_2.to(dtype=dtype)
 
 
 def load_caption_model():