Update app.py

app.py

@@ -53,64 +53,64 @@ 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 = AutoTokenizer.from_pretrained(tokenizer_path)
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
 
-# text_encoder = WanT5EncoderModel.from_pretrained(
-#     text_encoder_path,
-#     additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-#     low_cpu_mem_usage=False,
-#     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')
-# transformer3d = WanTransformer3DModel.from_pretrained(
-#     transformer_path,
-#     subfolder=transformer_subpath,
-#     transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
-# )
-
-# noise_scheduler = FlowMatchEulerDiscreteScheduler(
-#     **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
-# )
-
-
-tokenizer = AutoTokenizer.from_pretrained(
-    os.path.join(pretrained_model_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
-)
 text_encoder = WanT5EncoderModel.from_pretrained(
-    os.path.join(pretrained_model_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+    text_encoder_path,
     additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
-    low_cpu_mem_usage=True,
-)
-clip_image_encoder = CLIPModel.from_pretrained(
-    os.path.join(pretrained_model_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
+    low_cpu_mem_usage=False,
+    torch_dtype=torch.bfloat16
 )
+
+clip_image_encoder = CLIPModel.from_pretrained(image_encoder_path)
+
 vae = AutoencoderKLWan.from_pretrained(
-    os.path.join(pretrained_model_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
+    vae_path,
     additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
 )
+
+transformer_subpath = config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')
 transformer3d = WanTransformer3DModel.from_pretrained(
-    os.path.join(transformer_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+    transformer_path,
+    subfolder=transformer_subpath,
     transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
 )
+
 noise_scheduler = FlowMatchEulerDiscreteScheduler(
     **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
 )
 
+
+# tokenizer = AutoTokenizer.from_pretrained(
+#     os.path.join(pretrained_model_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
+# )
+# text_encoder = WanT5EncoderModel.from_pretrained(
+#     os.path.join(pretrained_model_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+#     additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+#     low_cpu_mem_usage=True,
+# )
+# clip_image_encoder = CLIPModel.from_pretrained(
+#     os.path.join(pretrained_model_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
+# )
+# vae = AutoencoderKLWan.from_pretrained(
+#     os.path.join(pretrained_model_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
+#     additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+# )
+# transformer3d = WanTransformer3DModel.from_pretrained(
+#     os.path.join(transformer_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+#     transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+# )
+# noise_scheduler = FlowMatchEulerDiscreteScheduler(
+#     **filter_kwargs(FlowMatchEulerDiscreteScheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+# )
+
 pipeline = WanFunInpaintPipeline(
     vae=vae,
     text_encoder=text_encoder,