Update app.py

app.py

@@ -14,25 +14,26 @@ from diffsynth import ModelManager, WanVideoPipeline, save_video
 
 
 num_frames, width, height = 49, 832, 480
-gpu_id = 0
-device = f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu'
+# gpu_id = 3
+# device = f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu'
+device = f'cuda' if torch.cuda.is_available() else 'cpu'
 # pip install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu124
 
-# from modelscope import snapshot_download
-# model_dir = snapshot_download( # https://www.modelscope.cn/models/AI-ModelScope/RMBG-2.0
-#     model_id = 'AI-ModelScope/RMBG-2.0',
-#     local_dir = 'ckpt/RMBG-2.0',
-#     ignore_file_pattern = ['onnx*'],
-# )
+from modelscope import snapshot_download
+model_dir = snapshot_download( # https://www.modelscope.cn/models/AI-ModelScope/RMBG-2.0
+    model_id = 'AI-ModelScope/RMBG-2.0',
+    local_dir = 'ckpt/RMBG-2.0',
+    ignore_file_pattern = ['onnx*'],
+)
 
-# from huggingface_hub import snapshot_download, hf_hub_download
-# 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,
-#     resume_download=True,
-#     repo_type="model"
-# )
+from huggingface_hub import snapshot_download, hf_hub_download
+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,
+    resume_download=True,
+    repo_type="model"
+)
 
 # hf_hub_download(
 #     repo_id="Kunbyte/Lumen",
@@ -42,37 +43,37 @@ device = f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu'
 #     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()
+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
-# 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'
+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(
-#         [
-#             wan_dit_path if wan_dit_path else 'ckpt/Wan2.1-Fun-14B-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', 
-#         ], 
-#         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', 
-#             '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', 
-#         ], 
-#         torch_dtype=torch.bfloat16, 
-#     )
-# wan_pipe = WanVideoPipeline.from_model_manager(model_manager, torch_dtype=torch.bfloat16, device=device)
-# wan_pipe.enable_vram_management(num_persistent_param_in_dit=None)
+if 'wan14b' in wan_dit_path.lower(): # 14B: uses about 36G, about 10 min per video 
+    model_manager.load_models(
+        [
+            wan_dit_path if wan_dit_path else 'ckpt/Wan2.1-Fun-14B-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', 
+        ], 
+        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', 
+            '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', 
+        ], 
+        torch_dtype=torch.bfloat16, 
+    )
+wan_pipe = WanVideoPipeline.from_model_manager(model_manager, torch_dtype=torch.bfloat16, device=device)
+wan_pipe.enable_vram_management(num_persistent_param_in_dit=None)
 
 gr_info_duration = 2 # gradio popup information duration
 
@@ -196,7 +197,7 @@ video_dir = 'test/pachong_test/video/single'
 relight_dir = ''
 
 header = """
-# 💡Lumen: Consistent Video Relighting and Harmonious Background Replacement\n # <center>with Video Generative Models </center>
+# <center>💡Lumen: Consistent Video Relighting and Harmonious Background Replacement with Video Generative Models </center>
 
 <div style="text-align: center; display: flex; justify-content: left; gap: 5px;">
 <a href="https://lumen-relight.github.io"><img src="https://img.shields.io/badge/Project%20Page-Lumen-blue" alt="Project"></a>
@@ -324,9 +325,4 @@ with gr.Blocks(title="Lumen: Video Relighting Model").queue() as demo:
 
 # Launch application
 if __name__ == "__main__":
-    demo.launch()
-    # demo.launch(
-    #     server_name='0.0.0.0', 
-    #     debug=True,
-    #     ssr_mode=False,
-    # )
+    demo.launch() # max_threads