Update app.py

app.py

@@ -15,15 +15,22 @@ num_frames, width, height = 49, 832, 480
 gpu_id = 0
 device = f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu'
 
-from huggingface_hub import snapshot_download, hf_hub_download
-# snapshot_download(
-#     repo_id="briaai/RMBG-2.0",
-#     local_dir="ckpt/RMBG-2.0",
-#     local_dir_use_symlinks=False,
-#     resume_download=True,
-#     repo_type="model" 
+# from modelscope import snapshot_download
+# model_dir = snapshot_download(
+#     model_id = 'AI-ModelScope/RMBG-2.0',
+#     local_dir = 'ckpt/',
 # )
-# snapshot_download( # 下载整个仓库
+
+from huggingface_hub import snapshot_download, hf_hub_download
+hf_hub_download(
+    repo_id="alibaba-pai/Wan2.1-Fun-1.3B-Control",
+    filename="Wan2.1_VAE.pth",
+    local_dir="ckpt/Wan2.1-Fun-1.3B-Control/",
+    local_dir_use_symlinks=False,
+    resume_download=True,
+)
+
+# snapshot_download( # 下载整个仓库; 下briaai/RMBG-2.0需要token
 #     repo_id="alibaba-pai/Wan2.1-Fun-1.3B-Control",
 #     local_dir="ckpt/Wan2.1-Fun-1.3B-Control",
 #     local_dir_use_symlinks=False,
@@ -31,14 +38,23 @@ from huggingface_hub import snapshot_download, hf_hub_download
 #     repo_type="model"
 # )
 
-# rmbg_model = AutoModelForImageSegmentation.from_pretrained('ckpt/RMBG-2.0', trust_remote_code=True)
+# hf_hub_download(
+#     repo_id="Kunbyte/Lumen",
+#     filename="Lumen-T2V-1.3B.ckpt",
+#     local_dir="ckpt/",
+#     local_dir_use_symlinks=False,
+#     resume_download=True,
+# )
+
+# rmbg_model = AutoModelForImageSegmentation.from_pretrained('ckpt/RMBG-2.0', trust_remote_code=True) # ckpt/RMBG-2.0
 # torch.set_float32_matmul_precision(['high', 'highest'][0])
 # rmbg_model.to(device)
 # rmbg_model.eval()
 
-model_manager = ModelManager(device="cpu") # 1.3b: device=cpu: uses 6G VRAM, device=device: uses 16G VRAM; about 1-2 min per video
 
+model_manager = ModelManager(device="cpu") # 1.3b: device=cpu: uses 6G VRAM, device=device: uses 16G VRAM; about 1-2 min per video
 wan_dit_path = 'train_res/wan1.3b_zh/full_wc0.5_f1gt0.5_real1_2_zh_en_l_s/lightning_logs/version_0/checkpoints/step-step=30000.ckpt'
+
 if 'wan14b' in wan_dit_path.lower(): # 14B: uses about 36G, about 10 min per video 
     model_manager.load_models(
         [
@@ -50,13 +66,12 @@ if 'wan14b' in wan_dit_path.lower(): # 14B: uses about 36G, about 10 min per vid
         torch_dtype=torch.bfloat16, # float8_e4m3fn fp8量化; bfloat16 
     ) 
 else:
-    wan_dit_path = None
     model_manager.load_models( 
         [ 
-            wan_dit_path if wan_dit_path else 'ckpt/Wan2.1-Fun-1.3B-Control/diffusion_pytorch_model.safetensors', 
+            # wan_dit_path if wan_dit_path else 'ckpt/Wan2.1-Fun-1.3B-Control/diffusion_pytorch_model.safetensors', 
             'ckpt/Wan2.1-Fun-1.3B-Control/Wan2.1_VAE.pth', 
-            'ckpt/Wan2.1-Fun-1.3B-Control/models_t5_umt5-xxl-enc-bf16.pth', 
-            'ckpt/Wan2.1-Fun-1.3B-Control/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth', 
+            # 'ckpt/Wan2.1-Fun-1.3B-Control/models_t5_umt5-xxl-enc-bf16.pth', 
+            # 'ckpt/Wan2.1-Fun-1.3B-Control/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth', 
         ], 
         torch_dtype=torch.bfloat16, 
     )