TheStageAI/Elastic-Wan2.2-T2V-A14B-Diffusers

Instalación y configuración:

pip install thestage
thestage config set --access-token

pip install thestage
pip install 'thestage-elastic-models[nvidia]' \
  --extra-index-url https://thestage.jfrog.io/artifactory/api/pypi/pypi-thestage-ai-production/simple
pip install diffusers==0.35.1
pip install flash_attn==2.7.3 --no-build-isolation
pip uninstall apex
pip install opencv-python==4.11.0.86 imageio-ffmpeg==0.6.0 imageio ftfy
pip install hf_transfer

Ejemplo de inferencia con la variante acelerada S:

import torch
from elastic_models.diffusers import WanPipeline
from diffusers.utils import export_to_video

model_name = 'Wan-AI/Wan2.2-T2V-A14B-Diffusers'
device = torch.device("cuda")
dtype = torch.bfloat16

pipe = WanPipeline.from_pretrained(
    model_name,
    torch_dtype=dtype,
    # 'original' for original model
    # 'S' for accelerated model
    mode='S'
)
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.to(device)

prompt = "A beautiful woman in a red dress dancing"

with torch.no_grad():
    output = pipe(
        prompt=prompt,
        negative_prompt="",
        height=480,
        width=480,
        num_frames=81,
        num_inference_steps=40,
        guidance_scale=3.0,
        guidance_scale_2=2.0,
        generator=torch.Generator("cuda").manual_seed(42),
    )

video = output.frames[0]
export_to_video(video, "output.mp4", fps=16)

Reproducción del benchmark de latencia:

import time
import torch
from elastic_models.diffusers import WanPipeline

model_name = 'Wan-AI/Wan2.2-T2V-A14B-Diffusers'
device = torch.device("cuda")

pipe = WanPipeline.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    mode='S'
)
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
pipe.to(device)

prompt = "A beautiful woman in a red dress dancing"
generate_kwargs = {
    "prompt": prompt,
    "negative_prompt": "",
    "height": 480,
    "width": 480,
    "num_frames": 81,
    "num_inference_steps": 40,
    "guidance_scale": 3.0,
    "guidance_scale_2": 2.0,
}

def evaluate_pipeline():
    torch.cuda.synchronize()
    start_time = time.time()
    with torch.no_grad():
        pipe(**generate_kwargs)
    torch.cuda.synchronize()
    return time.time() - start_time

# Warm-up
for _ in range(2):
    evaluate_pipeline()

# Benchmarking
num_runs = 5
total_time = sum(evaluate_pipeline() for _ in range(num_runs))
average_latency = total_time / num_runs
print(f"Average Latency over {num_runs} runs: {average_latency} seconds")