Update app.py

app.py

@@ -5,7 +5,7 @@ import spaces
 import numpy as np
 from PIL import Image
 import safetensors.torch
-from huggingface_hub import hf_hub_download
+from huggingface_hub import snapshot_download
 from accelerate import Accelerator
 from accelerate.utils import set_seed
 from diffusers import (
@@ -22,10 +22,11 @@ from myutils.wavelet_color_fix import wavelet_color_fix, adain_color_fix
 pipeline = None
 generator = None
 accelerator = None
+model_path = None
 
 @spaces.GPU
 def initialize_models():
-    global pipeline, generator, accelerator
+    global pipeline, generator, accelerator, model_path
     
     # Initialize accelerator
     accelerator = Accelerator(
@@ -34,50 +35,41 @@ def initialize_models():
     )
 
     try:
-        # Download and load models with authentication token
+        # Download the entire repository
+        model_path = snapshot_download(
+            repo_id="NightRaven109/CCSRModels",
+            token=os.environ['Read2']
+        )
+        
+        # Load models from local directory
         scheduler = DDPMScheduler.from_pretrained(
-            "NightRaven109/CCSRModels",
-            subfolder="stable-diffusion-2-1-base/scheduler",
-            use_auth_token=os.environ['Read2']
+            os.path.join(model_path, "stable-diffusion-2-1-base/scheduler")
         )
         
         text_encoder = CLIPTextModel.from_pretrained(
-            "NightRaven109/CCSRModels",
-            subfolder="stable-diffusion-2-1-base/text_encoder",
-            use_auth_token=os.environ['Read2']
+            os.path.join(model_path, "stable-diffusion-2-1-base/text_encoder")
         )
         
         tokenizer = CLIPTokenizer.from_pretrained(
-            "NightRaven109/CCSRModels",
-            subfolder="stable-diffusion-2-1-base/tokenizer",
-            use_auth_token=os.environ['Read2']
+            os.path.join(model_path, "stable-diffusion-2-1-base/tokenizer")
         )
         
         feature_extractor = CLIPImageProcessor.from_pretrained(
-            "NightRaven109/CCSRModels",
-            subfolder="stable-diffusion-2-1-base/feature_extractor",
-            use_auth_token=os.environ['Read2']
+            os.path.join(model_path, "stable-diffusion-2-1-base/feature_extractor")
         )
         
         unet = UNet2DConditionModel.from_pretrained(
-            "NightRaven109/CCSRModels",
-            subfolder="stable-diffusion-2-1-base/unet",
-            use_auth_token=os.environ['Read2']
+            os.path.join(model_path, "stable-diffusion-2-1-base/unet")
         )
         
         controlnet = ControlNetModel.from_pretrained(
-            "NightRaven109/CCSRModels",
-            subfolder="Controlnet",
-            use_auth_token=os.environ['Read2']
+            os.path.join(model_path, "Controlnet")
         )
         
         vae = AutoencoderKL.from_pretrained(
-            "NightRaven109/CCSRModels",
-            subfolder="vae",
-            use_auth_token=os.environ['Read2']
+            os.path.join(model_path, "vae")
         )
 
-        # Rest of the code remains the same
         # Freeze models
         for model in [vae, text_encoder, unet, controlnet]:
             model.requires_grad_(False)