CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node 2 infer_helios.py \
    --enable_parallelism \
    --base_model_path "BestWishYsh/Helios-Base" \
    --transformer_path "BestWishYsh/Helios-Base" \
    --sample_type "t2v" \
    --num_frames 99 \
    --fps 24 \
    --prompt "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement." \
    --guidance_scale 5.0 \
    --output_folder "./output_helios/helios-base"

import torch
from diffusers import ModularPipeline, ClassifierFreeGuidance
from diffusers.utils import export_to_video, load_image, load_video

mod_pipe = ModularPipeline.from_pretrained("BestWishYsh/Helios-Distilled")
mod_pipe.load_components(torch_dtype=torch.bfloat16)
mod_pipe.to("cuda")

# we need to upload guider to the model repo, so each checkpoint will be able to config their guidance differently
guider = ClassifierFreeGuidance(guidance_scale=1.0)
mod_pipe.update_components(guider=guider)

# --- T2V ---
print("=== T2V ===")
prompt = (
    "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. "
    "The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving "
    "fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and "
    "sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef "
    "itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures "
    "the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. "
    "A close-up shot with dynamic movement."
)

output = mod_pipe(
    prompt=prompt,
    height=384,
    width=640,
    num_frames=240,
    pyramid_num_inference_steps_list=[2, 2, 2],
    is_amplify_first_chunk=True,
    generator=torch.Generator("cuda").manual_seed(42),
    output="videos",
)

export_to_video(output[0], "helios_distilled_modular_t2v_output.mp4", fps=24)
print(f"T2V max memory: {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB")
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

# --- I2V ---
print("=== I2V ===")
image = load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/wave.jpg"
)
i2v_prompt = (
    "A towering emerald wave surges forward, its crest curling with raw power and energy. "
    "Sunlight glints off the translucent water, illuminating the intricate textures and deep green hues within the wave's body."
)

output = mod_pipe(
    prompt=i2v_prompt,
    image=image,
    height=384,
    width=640,
    num_frames=240,
    pyramid_num_inference_steps_list=[2, 2, 2],
    is_amplify_first_chunk=True,
    generator=torch.Generator("cuda").manual_seed(42),
    output="videos",
)

export_to_video(output[0], "helios_distilled_modular_i2v_output.mp4", fps=24)
print(f"I2V max memory: {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB")
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()

# --- V2V ---
print("=== V2V ===")
video = load_video(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/helios/car.mp4"
)
v2v_prompt = (
    "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train. "
    "The camera captures various elements such as lush green fields, towering trees, quaint countryside houses, "
    "and distant mountain ranges passing by quickly."
)

output = mod_pipe(
    prompt=v2v_prompt,
    video=video,
    height=384,
    width=640,
    num_frames=240,
    pyramid_num_inference_steps_list=[2, 2, 2],
    is_amplify_first_chunk=True,
    generator=torch.Generator("cuda").manual_seed(42),
    output="videos",
)

export_to_video(output[0], "helios_distilled_modular_v2v_output.mp4", fps=24)
print(f"V2V max memory: {torch.cuda.max_memory_allocated() / 1024**3:.3f} GB")

cd vllm-omni

# Helios-Base
python3 examples/offline_inference/helios/end2end.py \
  --sample-type t2v \
  --model ./Helios-Base \
  --prompt "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement." \
  --num-frames 99 \
  --seed 42 \
  --output helios_t2v_base.mp4

# Helios-Mid
python examples/offline_inference/helios/end2end.py \
  --model ./Helios-Mid --sample-type t2v \
  --prompt "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement." \
  --guidance-scale 5.0 --is-enable-stage2 \
  --pyramid-num-inference-steps-list 20 20 20 \
  --num-frames 99 \
  --use-cfg-zero-star --use-zero-init --zero-steps 1 \
  --output helios_t2v_mid.mp4

# Helios-Distilled
python examples/offline_inference/helios/end2end.py \
  --model ./Helios-Distilled --sample-type t2v \
  --prompt "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement." \
  --num-frames 240 --guidance-scale 1.0 --is-enable-stage2 \
  --pyramid-num-inference-steps-list 2 2 2 \
  --is-amplify-first-chunk --output helios_t2v_distilled.mp4

sglang generate \
  --model-path BestWishYsh/Helios-Base \
  --prompt "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement." \
  --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
  --height 384 \
  --width 640 \
  --num-frames 99 \
  --num-inference-steps 50 \
  --guidance-scale 5.0

sglang generate \
  --model-path BestWishYsh/Helios-Base \
  --prompt "A vibrant tropical fish swimming gracefully among colorful coral reefs in a clear, turquoise ocean. The fish has bright blue and yellow scales with a small, distinctive orange spot on its side, its fins moving fluidly. The coral reefs are alive with a variety of marine life, including small schools of colorful fish and sea turtles gliding by. The water is crystal clear, allowing for a view of the sandy ocean floor below. The reef itself is adorned with a mix of hard and soft corals in shades of red, orange, and green. The photo captures the fish from a slightly elevated angle, emphasizing its lively movements and the vivid colors of its surroundings. A close-up shot with dynamic movement." \
  --negative-prompt "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" \
  --height 384 \
  --width 640 \
  --num-frames 99 \
  --num-inference-steps 50 \
  --guidance-scale 5.0 \
  --backend diffusers