phughk/krea-zill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Working script?

Browse Source
This commit is contained in:
Hugh Kaznowski
2026-09-10 22:23:26 +01:00
parent 6713dd598f
commit 1a488b7b5f
3 changed files with 1994 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.py
+197 -1
View File
@@ -1,5 +1,201 @@
"""Simple text-to-video interface for krea/krea-realtime-video, tuned for Apple Silicon (MPS).
This is *not* the realtime/streaming server from the model card (that one assumes CUDA,
flash_attn, and a B200 GPU). Instead it drives the diffusers ModularPipeline directly,
running the standard (non-realtime) block-by-block denoising loop on MPS, and exposes
the generation knobs the pipeline actually supports through a small Gradio UI.
Run with:
uv run main.py
"""
import os
import time
import gradio as gr
import torch
from diffusers import ModularPipeline
from diffusers.modular_pipelines import PipelineState
from diffusers.utils import export_to_video
REPO_ID = "krea/krea-realtime-video"
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "outputs")
DTYPE_CHOICES = {
"bfloat16 (recommended on M3/M4)": torch.bfloat16,
"float16 (safer on M1/M2)": torch.float16,
"float32 (slow, most compatible)": torch.float32,
}
_pipe = None
def get_device() -> str:
if torch.backends.mps.is_available():
return "mps"
print("MPS not available, falling back to CPU (this will be very slow).")
return "cpu"
def load_pipeline(dtype: torch.dtype, progress: gr.Progress | None = None):
"""Load (and cache) the modular pipeline. Downloads ~28GB of weights on first run."""
global _pipe
if _pipe is not None and _pipe._dtype == dtype:
return _pipe
if progress is not None:
progress(0, desc="Loading pipeline (first run downloads ~28GB from Hugging Face)...")
device = get_device()
pipe = ModularPipeline.from_pretrained(REPO_ID, trust_remote_code=True)
pipe.load_components(
trust_remote_code=True,
device_map=device,
torch_dtype={"default": dtype, "vae": torch.float32},
)
for block in pipe.transformer.blocks:
block.self_attn.fuse_projections()
pipe._dtype = dtype
_pipe = pipe
return pipe
def generate_video(
prompt: str,
negative_prompt: str,
height: int,
width: int,
num_inference_steps: int,
num_blocks: int,
num_frames_per_block: int,
fps: int,
seed: int,
dtype_label: str,
lora_repo: str,
lora_weight_name: str,
lora_adapter_name: str,
progress=gr.Progress(),
):
if not prompt or not prompt.strip():
raise gr.Error("Prompt cannot be empty.")
dtype = DTYPE_CHOICES[dtype_label]
pipe = load_pipeline(dtype, progress)
if lora_repo.strip():
progress(0, desc=f"Loading LoRA from {lora_repo}...")
pipe.transformer.load_lora_adapter(
lora_repo.strip(),
prefix="diffusion_model",
weight_name=lora_weight_name.strip() or None,
adapter_name=(lora_adapter_name.strip() or "user_lora"),
)
device = get_device()
seed = int(seed)
if seed < 0:
seed = torch.seed() % (2**32)
generator = torch.Generator(device=device if device != "mps" else "cpu").manual_seed(seed)
frames = []
state = PipelineState()
start = time.time()
for block_idx in range(num_blocks):
progress(block_idx / num_blocks, desc=f"Denoising block {block_idx + 1}/{num_blocks}...")
state = pipe(
state,
prompt=[prompt],
negative_prompt=[negative_prompt] if negative_prompt.strip() else None,
height=height,
width=width,
num_inference_steps=num_inference_steps,
num_blocks=num_blocks,
num_frames_per_block=num_frames_per_block,
block_idx=block_idx,
generator=generator,
)
frames.extend(state.values["videos"][0])
elapsed = time.time() - start
progress(1.0, desc="Exporting video...")
os.makedirs(OUTPUT_DIR, exist_ok=True)
out_path = os.path.join(OUTPUT_DIR, f"t2v_seed{seed}_{int(time.time())}.mp4")
export_to_video(frames, out_path, fps=fps)
info = f"Seed: {seed} | {len(frames)} frames @ {fps}fps | generated in {elapsed:.1f}s on {device}"
return out_path, info
with gr.Blocks(title="Krea Realtime Video - Text to Video") as demo:
gr.Markdown(
"# Krea Realtime Video — Text to Video\n"
"Local text-to-video with `krea/krea-realtime-video` (Wan2.1 14B distilled) via 🧨 diffusers, "
"running on Apple Silicon (MPS). First generation downloads ~28GB of weights."
)
with gr.Row():
with gr.Column(scale=1):
prompt = gr.Textbox(label="Prompt", placeholder="a cat sitting on a boat", lines=3)
negative_prompt = gr.Textbox(label="Negative prompt (optional)", lines=2)
with gr.Row():
height = gr.Slider(256, 720, value=480, step=16, label="Height")
width = gr.Slider(256, 1280, value=832, step=16, label="Width")
with gr.Row():
num_inference_steps = gr.Slider(1, 12, value=6, step=1, label="Inference steps per block")
fps = gr.Slider(8, 30, value=24, step=1, label="Output FPS")
with gr.Row():
num_blocks = gr.Slider(1, 30, value=9, step=1, label="Number of blocks (length)")
num_frames_per_block = gr.Slider(1, 12, value=3, step=1, label="Frames per block")
with gr.Row():
seed = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
dtype_label = gr.Dropdown(
choices=list(DTYPE_CHOICES.keys()),
value="bfloat16 (recommended on M3/M4)",
label="Compute dtype",
)
with gr.Accordion("LoRA (optional)", open=False):
lora_repo = gr.Textbox(label="LoRA repo id", placeholder="shauray/Origami_WanLora")
lora_weight_name = gr.Textbox(
label="LoRA weight filename", placeholder="origami_000000500.safetensors"
)
lora_adapter_name = gr.Textbox(label="Adapter name", placeholder="origami")
generate_btn = gr.Button("Generate", variant="primary")
with gr.Column(scale=1):
video_out = gr.Video(label="Result")
info_out = gr.Textbox(label="Run info", interactive=False)
generate_btn.click(
fn=generate_video,
inputs=[
prompt,
negative_prompt,
height,
width,
num_inference_steps,
num_blocks,
num_frames_per_block,
fps,
seed,
dtype_label,
lora_repo,
lora_weight_name,
lora_adapter_name,
],
outputs=[video_out, info_out],
)
def main():
print("Hello from krea-zill!")
demo.queue().launch()
if __name__ == "__main__":