Upload app.py with huggingface_hub

app.py

@@ -85,7 +85,8 @@ def load_models():
     """Load all models once and cache them globally."""
     global vae, tokenizer, text_encoder, unet, scheduler, device
     
-    if vae is not None:
+    # Check if already loaded
+    if vae is not None and scheduler is not None:
         return
     
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -93,20 +94,30 @@ def load_models():
     
     model_id = "CompVis/stable-diffusion-v1-4"
     
-    print("Loading models...")
-    vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae").to(device)
-    tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
-    text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder").to(device)
-    unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet").to(device)
-    
-    scheduler = LMSDiscreteScheduler(
-        beta_start=0.00085,
-        beta_end=0.012,
-        beta_schedule="scaled_linear",
-        num_train_timesteps=1000
-    )
-    
-    print("Models loaded successfully!")
+    try:
+        print("Loading models... (this may take a few minutes on CPU)")
+        
+        # Load with float16 on GPU, float32 on CPU
+        dtype = torch.float16 if device == "cuda" else torch.float32
+        
+        vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae", torch_dtype=dtype).to(device)
+        tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
+        text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder", torch_dtype=dtype).to(device)
+        unet = UNet2DConditionModel.from_pretrained(model_id, subfolder="unet", torch_dtype=dtype).to(device)
+        
+        # Initialize scheduler
+        scheduler = LMSDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000
+        )
+        
+        print("Models loaded successfully!")
+        
+    except Exception as e:
+        print(f"Error loading models: {e}")
+        raise RuntimeError(f"Failed to load models: {e}")
 
 
 def generate_with_style(