103 lines
3.3 KiB
Python
103 lines
3.3 KiB
Python
"""Smoke test: single DIT denoising step with GGUF weights.
|
|
|
|
Builds the pipeline, runs all encoders, initializes the scheduler, then
|
|
executes exactly one DIT forward pass (step_pre → infer → step_post).
|
|
This isolates the GGUF fp16 DIT from the rest of the pipeline.
|
|
|
|
Run:
|
|
docker compose exec -e DIT_QUANT=gguf-Q4_K_M voice-chat \
|
|
python -m tests.component.test_12_dit_single_step
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
|
|
from tests.component._common import ensure_sample_avatar, get_logger
|
|
|
|
log = get_logger("test_12")
|
|
|
|
DIT_QUANT = os.environ.get("DIT_QUANT", "gguf-Q8_0")
|
|
CONFIG_JSON = "/app/configs/lightx2v/wan22_i2v_gguf_5b_turbo.json"
|
|
DIT_REPO = "hum-ma/Wan2.2-TI2V-5B-Turbo-GGUF"
|
|
|
|
|
|
def run():
|
|
try:
|
|
from server.video_models.wan22 import Wan22Pipeline
|
|
except ImportError as e:
|
|
log.error("Import failed: %s", e)
|
|
sys.exit(0)
|
|
|
|
avatar = ensure_sample_avatar()
|
|
log.info("Avatar: %s", avatar)
|
|
|
|
log.info("Building pipeline (quant=%s)...", DIT_QUANT)
|
|
pipe = Wan22Pipeline(
|
|
base_repo="Wan-AI/Wan2.2-TI2V-5B",
|
|
dit_repo=DIT_REPO,
|
|
config_json=CONFIG_JSON,
|
|
model_cls="wan2.2",
|
|
resolution=480,
|
|
fps=16,
|
|
dit_quant_scheme=DIT_QUANT,
|
|
t5_quantized=True,
|
|
)
|
|
log.info("Pipeline ready.")
|
|
|
|
runner = pipe._runner
|
|
|
|
# Set up input_info for a short clip
|
|
from lightx2v.utils.input_info import update_input_info_from_dict
|
|
update_input_info_from_dict(
|
|
pipe._input_info_template,
|
|
{
|
|
"seed": 42,
|
|
"prompt": "a person looking at the camera, natural lighting",
|
|
"negative_prompt": "",
|
|
"image_path": avatar,
|
|
"target_video_length": 17, # 1 second at 16fps + 1
|
|
},
|
|
)
|
|
runner.input_info = pipe._input_info_template
|
|
|
|
# 1. Run all encoders (T5 + CLIP + VAE)
|
|
log.info("Running all input encoders (T5 + CLIP + VAE)...")
|
|
runner.inputs = runner.run_input_encoder()
|
|
log.info("Encoder outputs ready.")
|
|
for k, v in runner.inputs.items():
|
|
if isinstance(v, torch.Tensor):
|
|
log.info(" inputs[%s]: shape=%s dtype=%s", k, v.shape, v.dtype)
|
|
elif isinstance(v, dict):
|
|
for k2, v2 in v.items():
|
|
if isinstance(v2, torch.Tensor):
|
|
log.info(" inputs[%s][%s]: shape=%s dtype=%s", k, k2, v2.shape, v2.dtype)
|
|
|
|
# 2. Initialize run (sets up scheduler, creates noise latents)
|
|
log.info("Initializing run (scheduler.prepare)...")
|
|
runner.init_run()
|
|
latents = runner.model.scheduler.latents
|
|
log.info("Initial latents: shape=%s dtype=%s", latents.shape, latents.dtype)
|
|
|
|
# 3. Single DIT step
|
|
log.info("Running single DIT step (step_pre → infer → step_post)...")
|
|
runner.model.scheduler.step_pre(step_index=0)
|
|
runner.model.infer(runner.inputs)
|
|
runner.model.scheduler.step_post()
|
|
|
|
latents_after = runner.model.scheduler.latents
|
|
log.info("Latents after step: shape=%s dtype=%s", latents_after.shape, latents_after.dtype)
|
|
|
|
# Verify latents changed (denoising did something)
|
|
assert not torch.equal(latents, latents_after), "Latents unchanged after DIT step"
|
|
log.info("PASS: DIT single step completed, latents updated.")
|
|
|
|
log.info("PASS: DIT forward pass succeeded under %s pipeline.", DIT_QUANT)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
run()
|