rename

app.py

@@ -1,4 +1,6 @@
 import time
+import os
+from huggingface_hub import hf_hub_download
 
 import gradio as gr
 import spaces
@@ -18,7 +20,7 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
 def _load_pipeline_internal(
-    pretrained_model_path="THUDM/CogVideoX-5b",
+    pretrained_model_path="LiuhanChen/Sci-Fi",
     ef_net_path="weights/EF_Net.pth",
     dtype_str="bfloat16",
 ):
@@ -31,19 +33,34 @@ def _load_pipeline_internal(
 
     dtype = torch.float16 if dtype_str == "float16" else torch.bfloat16
 
-    # Load models
+    # Download EF-Net weights if not exists
+    if not os.path.exists(ef_net_path):
+        print("Downloading EF-Net weights from Hugging Face...")
+        os.makedirs("weights", exist_ok=True)
+        ef_net_path = hf_hub_download(
+            repo_id="LiuhanChen/Sci-Fi",
+            subfolder="EF_Net",
+            filename="EF_Net.pth",
+            local_dir="weights"
+        )
+        ef_net_path = "weights/EF_Net/EF_Net.pth"
+        print(f"EF-Net weights downloaded to {ef_net_path}")
+
+    # Load models from Hugging Face
     tokenizer = T5Tokenizer.from_pretrained(
-        pretrained_model_path, subfolder="tokenizer"
+        pretrained_model_path, subfolder="CogVideoX-5b-I2V/tokenizer"
     )
     text_encoder = T5EncoderModel.from_pretrained(
-        pretrained_model_path, subfolder="text_encoder"
+        pretrained_model_path, subfolder="CogVideoX-5b-I2V/text_encoder"
     )
     transformer = CustomCogVideoXTransformer3DModel.from_pretrained(
-        pretrained_model_path, subfolder="transformer"
+        pretrained_model_path, subfolder="CogVideoX-5b-I2V/transformer"
+    )
+    vae = AutoencoderKLCogVideoX.from_pretrained(
+        pretrained_model_path, subfolder="CogVideoX-5b-I2V/vae"
     )
-    vae = AutoencoderKLCogVideoX.from_pretrained(pretrained_model_path, subfolder="vae")
     scheduler = CogVideoXDDIMScheduler.from_pretrained(
-        pretrained_model_path, subfolder="scheduler"
+        pretrained_model_path, subfolder="CogVideoX-5b-I2V/scheduler"
     )
 
     # Load EF-Net