from fastvideo import VideoGenerator

OUTPUT_PATH = "video_samples"


def main():
    # FastVideo will automatically use the optimal default arguments for the
    # model.
    # If a local path is provided, FastVideo will make a best effort
    # attempt to identify the optimal arguments.
    generator = VideoGenerator.from_pretrained(
        "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
        # FastVideo will automatically handle distributed setup
        {
            "engine": {
                "num_gpus": 1,
                "use_fsdp_inference": False,  # set to True if GPU is out of memory
                "offload": {
                    "dit": False,
                    "vae": False,
                    "text_encoder": True,
                    "pin_cpu_memory": True,  # set to false if low CPU RAM or hit obscure "CUDA error: Invalid argument"
                    # "image_encoder": False,
                },
            },
        },
    )

    # Generate videos with the same simple API, regardless of GPU count
    prompt = ("A curious raccoon peers through a vibrant field of yellow sunflowers, its eyes "
              "wide with interest. The playful yet serene atmosphere is complemented by soft "
              "natural light filtering through the petals. Mid-shot, warm and cheerful tones.")
    video = generator.generate({
        "prompt": prompt,
        "output": {"output_path": OUTPUT_PATH, "save_video": True},
    })
    # video = generator.generate({
    #     "prompt": prompt,
    #     "inputs": {"image_path": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg"},
    #     "sampling": {"num_frames": 45},
    #     "output": {"output_path": "wan_t2v_videos/"},
    # })

    # Generate another video with a different prompt, without reloading the
    # model!
    prompt2 = ("A majestic lion strides across the golden savanna, its powerful frame "
               "glistening under the warm afternoon sun. The tall grass ripples gently in "
               "the breeze, enhancing the lion's commanding presence. The tone is vibrant, "
               "embodying the raw energy of the wild. Low angle, steady tracking shot, "
               "cinematic.")
    video2 = generator.generate({
        "prompt": prompt2,
        "output": {"output_path": OUTPUT_PATH, "save_video": True},
    })


if __name__ == "__main__":
    main()