Update app.py

app.py

@@ -38,15 +38,24 @@ def filter_kwargs(cls, kwargs):
 
 from huggingface_hub import snapshot_download
 
-def download_component_subfolder(repo_id, subfolder):
-    local_dir = snapshot_download(
+# def download_component_subfolder(repo_id, subfolder):
+#     local_dir = snapshot_download(
+#         repo_id=repo_id,
+#         repo_type="model",
+#         local_dir="ckpt/Wan2.1-Fun-1.3B-InP",
+#         local_dir_use_symlinks=False,
+#         # allow_patterns=[f"{subfolder}/*"]
+#     )
+#     return os.path.join(local_dir, subfolder)
+
+def download_model_component(repo_id, subpath):
+    return os.path.join(snapshot_download(
         repo_id=repo_id,
         repo_type="model",
-        local_dir="ckpt/Wan2.1-Fun-1.3B-InP",
+        local_dir=f"ckpt/{repo_id.replace('/', '--')}",
         local_dir_use_symlinks=False,
-        # allow_patterns=[f"{subfolder}/*"]
-    )
-    return os.path.join(local_dir, subfolder)
+        allow_patterns=[f"{subpath}/*"]
+    ), subpath)
 
 # pretrained_model_path = "./models/Diffusion_Transformer/Wan2.1-Fun-1.3B-InP"
 pretrained_model_path = "alibaba-pai/Wan2.1-Fun-1.3B-InP"
@@ -55,12 +64,17 @@ transformer_path = "Kunbyte/ROSE"
 config_path = "configs/wan2.1/wan_civitai.yaml"
 config = OmegaConf.load(config_path)
 
-repo_id = "alibaba-pai/Wan2.1-Fun-1.3B-InP"
+# repo_id = "alibaba-pai/Wan2.1-Fun-1.3B-InP"
 
-text_encoder_path = download_component_subfolder(repo_id, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder'))
-tokenizer_path = download_component_subfolder(repo_id, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer'))
-image_encoder_path = download_component_subfolder(repo_id, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder'))
-vae_path = download_component_subfolder(repo_id, config['vae_kwargs'].get('vae_subpath', 'vae'))
+# text_encoder_path = download_component_subfolder(repo_id, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder'))
+# tokenizer_path = download_component_subfolder(repo_id, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer'))
+# image_encoder_path = download_component_subfolder(repo_id, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder'))
+# vae_path = download_component_subfolder(repo_id, config['vae_kwargs'].get('vae_subpath', 'vae'))
+tokenizer_path = download_model_component("alibaba-pai/Wan2.1-Fun-1.3B-InP", config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer'))
+text_encoder_path = download_model_component("alibaba-pai/Wan2.1-Fun-1.3B-InP", config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder'))
+image_encoder_path = download_model_component("alibaba-pai/Wan2.1-Fun-1.3B-InP", config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder'))
+vae_path = download_model_component("alibaba-pai/Wan2.1-Fun-1.3B-InP", config['vae_kwargs'].get('vae_subpath', 'vae'))
+transformer_path = download_model_component("Kunbyte/ROSE", config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer'))
 
 tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
 
@@ -81,7 +95,6 @@ vae = AutoencoderKLWan.from_pretrained(
 transformer_subpath = config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')
 transformer3d = WanTransformer3DModel.from_pretrained(
     transformer_path,
-    subfolder=transformer_subpath,
     transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
 )