get vae in fp32 when using wan.

Author: sayakpaul

README.md

@@ -14,6 +14,8 @@ But it's been growing now! Check out the rest of the README to know more 🤗
 
 **Updates**
 
+🔥 04/03/2025: Support for LTX-Video and Wan in [this PR](https://github.com/sayakpaul/tt-scale-flux/pull/18) 🎬 Check out this section for results and more info.
+
 🔥 01/03/2025: `OpenAIVerifier` was added in [this PR](https://github.com/sayakpaul/tt-scale-flux/pull/16). Specify "openai" in the `name` under `verifier_args`. Thanks to [zhuole1025](https://github.com/zhuole1025) for contributing this!
 
 🔥 27/02/2025: [MaximClouser](https://github.com/MaximClouser) implemented a ComfyUI node for inference-time

main.py

@@ -221,7 +221,14 @@ def main():
 
     # === Set up the image-generation pipeline ===
     torch_dtype = TORCH_DTYPE_MAP[config.pop("torch_dtype")]
-    pipe = DiffusionPipeline.from_pretrained(pipeline_name, torch_dtype=torch_dtype)
+    fp_kwargs = {"pretrained_model_name_or_path": pipeline_name, "torch_dtype": torch_dtype}
+    if "Wan" in pipeline_name:
+        # As per recommendations from https://huggingface.co/docs/diffusers/main/en/api/pipelines/wan.
+        from diffusers import AutoencoderKLWan
+
+        vae = AutoencoderKLWan.from_pretrained(pipeline_name, subfolder="vae", torch_dtype=torch.float32)
+        fp_kwargs.update({"vae": vae})
+    pipe = DiffusionPipeline.from_pretrained(**fp_kwargs)
     if not config.get("use_low_gpu_vram", False):
         pipe = pipe.to("cuda:0")
     pipe.set_progress_bar_config(disable=True)