Update app.py

app.py

@@ -36,6 +36,16 @@ def filter_kwargs(cls, kwargs):
     valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
     return {k: v for k, v in kwargs.items() if k in valid_params}
 
+from huggingface_hub import snapshot_download
+
+def download_component_subfolder(repo_id, subfolder):
+    local_dir = snapshot_download(
+        repo_id=repo_id,
+        repo_type="model",
+        allow_patterns=[f"{subfolder}/*"]
+    )
+    return os.path.join(local_dir, subfolder)
+
 # pretrained_model_path = "./models/Diffusion_Transformer/Wan2.1-Fun-1.3B-InP"
 pretrained_model_path = "alibaba-pai/Wan2.1-Fun-1.3B-InP"
 transformer_path = "Kunbyte/ROSE"
@@ -43,12 +53,27 @@ transformer_path = "Kunbyte/ROSE"
 config_path = "./configs/wan2.1/wan_civitai.yaml"
 config = OmegaConf.load(config_path)
 
-text_encoder_subpath = config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')
+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'))
+
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
+
 text_encoder = WanT5EncoderModel.from_pretrained(
-    pretrained_model_path,
-    subfolder=text_encoder_subpath,
+    text_encoder_path,
     additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
     low_cpu_mem_usage=True,
+    torch_dtype=torch.bfloat16
+)
+
+clip_image_encoder = CLIPModel.from_pretrained(image_encoder_path)
+
+vae = AutoencoderKLWan.from_pretrained(
+    vae_path,
+    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
 )
 
 transformer_subpath = config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')
@@ -58,22 +83,11 @@ transformer3d = WanTransformer3DModel.from_pretrained(
     transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
 )
 
-image_encoder_subpath = config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')
-clip_image_encoder = CLIPModel.from_pretrained(pretrained_model_path, subfolder=image_encoder_subpath)
-
-vae_subpath = config['vae_kwargs'].get('vae_subpath', 'vae')
-vae = AutoencoderKLWan.from_pretrained(
-    pretrained_model_path,
-    subfolder=vae_subpath,
-    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
-)
-
 noise_scheduler = FlowMatchEulerDiscreteScheduler(
     **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
 )
 
-tokenizer_subpath = config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')
-tokenizer = AutoTokenizer.from_pretrained(pretrained_model_path, tokenizer_subpath)
+
 # tokenizer = AutoTokenizer.from_pretrained(
 #     os.path.join(pretrained_model_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
 # )