| from causvid.models.wan.wan_wrapper import CausalWanDiffusionWrapper, WanTextEncoder, WanVAEWrapper | |
| from diffusers.utils import export_to_video | |
| import torch | |
| torch.set_grad_enabled(False) | |
| model = CausalWanDiffusionWrapper().to("cuda").to(torch.float32) | |
| encoder = WanTextEncoder().to("cuda").to(torch.float32) | |
| vae = WanVAEWrapper().to("cuda").to(torch.float32) | |
| model.set_module_grad( | |
| { | |
| "model": False | |
| } | |
| ) | |
| text_prompts = [r"""A stylish woman walks down a Tokyo street filled with warm glowing neon and animated city signage. She wears a black leather jacket, a long red dress, and black boots, and carries a black purse. She wears sunglasses and red lipstick. She walks confidently and casually. The street is damp and reflective, creating a mirror effect of the colorful lights. Many pedestrians walk about."""] * 2 | |
| conditional_dict = encoder( | |
| text_prompts=text_prompts | |
| ) | |
| noise = torch.randn( | |
| 2, 21, 16, 60, 104, generator=torch.Generator().manual_seed(42), dtype=torch.float32 | |
| ).to("cuda") | |
| timetep = 999 * torch.ones([2, 21], device="cuda", dtype=torch.long) | |
| with torch.no_grad(): | |
| output = model(noise, conditional_dict, timetep) | |
| video = vae.decode_to_pixel(output) | |
| video = (video * 0.5 + 0.5).cpu().detach().to(torch.float32)[0].permute(0, 2, 3, 1).numpy() | |
| export_to_video(video, "causal_output.mp4", fps=8) | |