Update app.py

app.py

@@ -10,16 +10,17 @@ from glob import glob
 from pathlib import Path
 from typing import Optional
 
-from diffusers import StableVideoDiffusionPipeline
+from diffusers import StableVideoDiffusionPipeline, UNetSpatioTemporalConditionControlNetModel
 from diffusers.utils import load_image, export_to_video
 
 import uuid
-# from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download
 
 
-# os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-# HF_TOKEN = os.environ.get("HF_TOKEN", None)
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+HF_TOKEN = os.environ.get("HF_TOKEN", None)
 # Constants
+base = "stabilityai/stable-video-diffusion-img2vid-xt"
 model = "ECNU-CILab/ExVideo-SVD-128f-v1"
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -37,11 +38,28 @@ JS = """function () {
   }
 }"""
 
+downloaded_model_path = hf_hub_download(
+        repo_id=model,
+        filename=model.fp16.safetensors,
+        local_dir="model"
+)
+
+MODEL_PATH = "./model/"
+
+
 
 # Ensure model and scheduler are initialized in GPU-enabled function
 if torch.cuda.is_available():
+
+    unet = UNetSpatioTemporalConditionControlNetModel.from_pretrained(
+        MODEL_PATH,
+        low_cpu_mem_usage=True,
+        variant="fp16",
+    )
+    
     pipe = StableVideoDiffusionPipeline.from_pretrained(
-    model, 
+    base,
+    unet=unet,
     torch_dtype=torch.float16, 
     variant="fp16").to("cuda")