Views
No views yet
Efficient-Large-Model/SANA-Video_2B_480p_diffusers| Checkpoint step | Filename | Preview |
|---|---|---|
| base (0) | base model | ![]() |
| 100 | lora_step_000100.pt | ![]() |
| 1,000 | lora_step_001000.pt | ![]() |
| 2,000 | lora_step_002000.pt | ![]() |
| 5,000 | lora_step_005000.pt | ![]() |
| 7,500 | lora_step_007500.pt | ![]() |
| 10,000 | lora_step_010000.pt | ![]() |
lora_step_*.pt).1from pathlib import Path
2import torch
3from diffusers import SanaVideoPipeline
4from diffusers.utils import export_to_video
5from peft import LoraConfig, get_peft_model
6from peft.utils import set_peft_model_state_dict
7
8# ---- Paths / config ----
9MODEL_ID = "Efficient-Large-Model/SANA-Video_2B_480p_diffusers"
10LORA_PATH = Path("lora_step_010000.pt") # example checkpoint
11
12DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
13DTYPE = torch.bfloat16 if torch.cuda.is_available() else torch.float32
14
15LORA_R = 16
16LORA_ALPHA = 32
17LORA_DROPOUT = 0.0
18LORA_TARGET_MODULES = ["proj_out", "to_q", "to_v", "to_k", "linear_2", "linear_1", "linear"]
19
20PROMPT = (
21 "A vintage slapstick 2D cartoon scene of a grey cat chasing a small brown mouse "
22 "in a colorful house, Tom and Jerry style, bold outlines, limited color palette, "
23 "exaggerated expressions, smooth character motion."
24)
25
26# ---- Load base pipeline ----
27pipe = SanaVideoPipeline.from_pretrained(
28 MODEL_ID,
29 torch_dtype=DTYPE,
30)
31
32pipe.vae.to(DEVICE, dtype=torch.float32) # VAE in fp32 is more stable
33
34# ---- Wrap transformer with LoRA ----
35lora_cfg = LoraConfig(
36 r=LORA_R,
37 lora_alpha=LORA_ALPHA,
38 lora_dropout=LORA_DROPOUT,
39 bias="none",
40 target_modules=LORA_TARGET_MODULES,
41)
42pipe.transformer = get_peft_model(pipe.transformer, lora_cfg)
43
44# ---- Load LoRA weights (handles torch.compile prefixes) ----
45state = torch.load(LORA_PATH, map_location="cpu")
46fixed_state = {}
47for k, v in state.items():
48 if k.startswith("_orig_mod."):
49 k = k[len("_orig_mod."):]
50 if k.startswith("module."):
51 k = k[len("module."):]
52 fixed_state[k] = v
53
54set_peft_model_state_dict(pipe.transformer, fixed_state)
55
56# ---- Move to device ----
57pipe.to(DEVICE)
58pipe.transformer.to(DEVICE, dtype=DTYPE)
59pipe.text_encoder.to(DEVICE, dtype=DTYPE)
60
61# ---- Inference ----
62with torch.no_grad():
63 out = pipe(
64 prompt=[PROMPT],
65 num_inference_steps=50,
66 guidance_scale=4.0,
67 height=224,
68 width=224,
69 use_resolution_binning=False,
70 )
71
72video_frames = out.frames[0] if hasattr(out, "frames") else out.videos[0]
73export_to_video(video_frames, "tomjerry_lora_sample.mp4", fps=16)
74print("Saved to tomjerry_lora_sample.mp4")r = 16alpha = 32dropout = 0.1proj_out, to_q, to_k, to_v, linear_1, linear_2, linear.AdamW / AdamW8bit (bitsandbytes).2e-4.8 videos (5s @ 16 FPS, 224×224).bf16 (torch.set_float32_matmul_precision("medium")).