Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Video VAE decode: bf16 autocast + batched tiles (probe stays fp32)
8c3d737 verified | from __future__ import annotations | |
| import os | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from common import MODEL_REPO, drop_cached, load_wire, snapshot | |
| NAME = "diffusers" | |
| LABEL = "diffusers (PR #2)" | |
| STATE = {} | |
| def load(partition: str) -> None: | |
| from diffusers.modular_pipelines.pika_h3_flex import PikaH3FlexBlocks | |
| trunk = "transformer_ref" if partition == "ref2va" else "transformer" | |
| started = time.time() | |
| blocks = PikaH3FlexBlocks() | |
| blocks.sub_blocks.pop("text_encoder") | |
| pipe = blocks.init_pipeline(MODEL_REPO) | |
| pipe.load_components(names=["vae", "audio_vae"], dtype=torch.float32) | |
| pipe.load_components(names=[trunk, "scheduler", "audio_scheduler"], dtype=torch.bfloat16) | |
| adapter = Path(snapshot([f"{partition}/*"])) / partition / "adapter.safetensors" | |
| pipe.load_lora_weights( | |
| str(adapter), | |
| **({"load_into_transformer_ref": True} if partition == "ref2va" else {}), | |
| ) | |
| drop_cached(adapter) | |
| getattr(pipe, trunk).requires_grad_(False).eval() | |
| attention = os.environ.get("PIKA_ATTENTION") | |
| if attention: | |
| getattr(pipe, trunk).set_attention_backend(attention) | |
| getattr(pipe, trunk).to("cuda") | |
| STATE.update(partition=partition, trunk=trunk, pipe=pipe, loaded_in=time.time() - started) | |
| print(f"[diffusers] {partition}: {trunk} + {partition} adapter ready in {STATE['loaded_in']:.0f}s", flush=True) | |
| def video_vae(): | |
| return STATE["pipe"].vae | |
| def transformer(): | |
| return getattr(STATE["pipe"], STATE["trunk"]) | |
| def prepare(task, prompt, first_image, last_image, references, height, width, num_frames): | |
| from diffusers.modular_pipelines.pika_h3_flex import ( | |
| PikaH3FlexAudioReference, | |
| PikaH3FlexImageReference, | |
| PikaH3FlexVideoReference, | |
| ) | |
| from PIL import Image | |
| kinds = {"image": PikaH3FlexImageReference, "video": PikaH3FlexVideoReference, "audio": PikaH3FlexAudioReference} | |
| inputs = {"height": height, "width": width, "num_frames": num_frames} | |
| if task in ("i2va", "fl2va") and first_image: | |
| inputs["image"] = Image.open(first_image).convert("RGB") | |
| if task in ("l2va", "fl2va") and last_image: | |
| inputs["last_image"] = Image.open(last_image).convert("RGB") | |
| if task == "ref2va": | |
| inputs["references"] = [kinds[kind].from_file(path) for kind, path in references] | |
| return inputs | |
| def run(prepared, wire_path, seed, progress=None, noise=None): | |
| pipe = STATE["pipe"] | |
| prompt_embeds, text_token_tags, _ = load_wire(wire_path) | |
| inputs = dict(prepared) | |
| if noise is not None: | |
| inputs["latents"], inputs["audio_latents"] = noise | |
| pipe.vae.to("cuda") | |
| pipe.audio_vae.to("cuda") | |
| if progress is not None: | |
| calls = {"n": 0} | |
| def hook(module, args, kwargs): | |
| calls["n"] += 1 | |
| progress(calls["n"] / 6, desc=f"Flex step {calls['n']} of 6") | |
| handle = transformer().register_forward_pre_hook(hook, with_kwargs=True) | |
| torch.cuda.synchronize() | |
| torch.cuda.reset_peak_memory_stats() | |
| start = time.perf_counter() | |
| try: | |
| with torch.inference_mode(): | |
| state = pipe( | |
| prompt_embeds=prompt_embeds.to("cuda"), | |
| text_token_tags=text_token_tags, | |
| generator=torch.Generator("cuda").manual_seed(int(seed)), | |
| output_type="np", | |
| **inputs, | |
| ) | |
| torch.cuda.synchronize() | |
| finally: | |
| if progress is not None: | |
| handle.remove() | |
| total = time.perf_counter() - start | |
| videos = state.get("videos") | |
| frames = np.asarray(videos[0]) | |
| if frames.dtype != np.uint8: | |
| frames = (np.clip(frames, 0.0, 1.0) * 255.0).round().astype(np.uint8) | |
| audio = state.get("audio")[0] | |
| wave = audio.float().cpu().numpy() if torch.is_tensor(audio) else np.asarray(audio, dtype=np.float32) | |
| timings = {"pipeline": total, "peak_gib": torch.cuda.max_memory_allocated() / 2**30} | |
| return frames, wave, int(state.get("sampling_rate")), timings, None | |