Spaces:
Running on Zero
Running on Zero
File size: 4,031 Bytes
4213590 8c3d737 4213590 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | 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
|