fixes: natural-speed/aspect, audio, progress, prompts, examples

app.py

@@ -8,10 +8,11 @@ import random
 import tempfile
 
 import numpy as np
+import imageio.v3 as iio
 import spaces
 import torch
 import gradio as gr
-from PIL import Image
+from PIL import Image, ImageOps
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
@@ -30,10 +31,7 @@ NUM_STEPS = len(DISTILLED_SIGMA_VALUES)  # 8-step distilled schedule
 MAX_SEED = np.iinfo(np.int32).max
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-RES_PRESETS = {
-    "Fast (768×448)": (768, 448),
-    "Quality (960×544)": (960, 544),
-}
+RES_PRESETS = {"Fast (768×448)": (768, 448), "Quality (960×544)": (960, 544)}
 FRAME_CHOICES = [49, 73, 97, 121]
 
 # --- Load pipeline once at module scope (ZeroGPU registers it) ---------------
@@ -47,48 +45,70 @@ pipe.set_adapters("colorize", LORA_SCALE)
 
 
 # --- Helpers ----------------------------------------------------------------
-def _resample(frames, n):
-    """Pick n evenly spaced frames (handles clips shorter or longer than n)."""
-    idx = np.linspace(0, len(frames) - 1, n).round().astype(int)
-    return [frames[i] for i in idx]
+def _src_fps(path, default=FPS):
+    try:
+        return float(iio.immeta(path, plugin="pyav").get("fps", default)) or default
+    except Exception:
+        return default
 
 
-def _to_grayscale(img: Image.Image) -> Image.Image:
-    """Desaturate to a monochrome reference, kept as 3-channel RGB."""
+def _load_frames(path, num_frames, width, height):
+    """Natural-speed (real-time at 24fps), aspect-preserving (center-crop) frames."""
+    frames = load_video(path)
+    if not frames:
+        return []
+    fps = _src_fps(path)
+    out = []
+    for i in range(num_frames):
+        idx = int(round(i / FPS * fps))
+        idx = min(idx, len(frames) - 1)
+        out.append(ImageOps.fit(frames[idx].convert("RGB"), (width, height), Image.LANCZOS))
+    return out
+
+
+def _to_grayscale(img):
     return img.convert("L").convert("RGB")
 
 
-def _pick_resolution(first_frame: Image.Image, preset: str):
+def _pick_resolution(first_frame, preset):
     w, h = RES_PRESETS[preset]
-    # Match the input's orientation (the model saw both landscape and portrait).
     if first_frame.height > first_frame.width:
         w, h = h, w
     return w, h
 
 
-def _build_prompt(colors: str, scene: str) -> str:
+def _build_prompt(colors, scene, audio):
     colors = colors.strip()
     scene = scene.strip() or "the same scene"
-    return (
+    p = (
         f"Reference shows {scene}, rendered in high-contrast monochrome with soft natural daylight. "
         f"Edited shows the same scene with natural colors restored. "
         f"COLORIZE {colors}. "
         f"Subject identity, framing, and background geometry are identical to the reference; "
         f"only color information differs between reference and edited."
     )
+    if audio.strip():
+        p += f" Audio: {audio.strip()}."
+    return p
+
+
+def _export(video_np, audio, path):
+    kw = {}
+    if audio is not None:
+        kw = dict(audio=audio[0].float().cpu(), audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
+    encode_video(video_np, fps=FPS, output_path=path, **kw)
 
 
 def _duration(*args, **kwargs):
-    # args mirror the GPU fn: (video, colors, scene, preset, num_frames, seed, randomize, [progress])
-    preset = args[3] if len(args) > 3 else "Fast"
-    num_frames = args[4] if len(args) > 4 else 73
+    preset = next((a for a in args if a in RES_PRESETS), "Fast")
+    num_frames = next((a for a in args if a in FRAME_CHOICES), 73)
     per_frame = 1.6 if "Quality" in str(preset) else 1.0
-    return int(50 + int(num_frames) * per_frame)
+    return int(60 + int(num_frames) * per_frame)
 
 
 # --- Inference --------------------------------------------------------------
 @spaces.GPU(duration=_duration)
-def colorize(video, colors, scene, preset, num_frames, seed, randomize,
+def colorize(video, colors, scene, audio, preset, num_frames, seed, randomize,
              progress=gr.Progress(track_tqdm=True)):
     if video is None:
         raise gr.Error("Please upload a video to colorize.")
@@ -98,43 +118,34 @@ def colorize(video, colors, scene, preset, num_frames, seed, randomize,
     if randomize:
         seed = random.randint(0, MAX_SEED)
     seed = int(seed)
+    num_frames = int(num_frames)
 
-    frames = load_video(video)
-    if not frames:
+    probe = load_video(video)
+    if not probe:
         raise gr.Error("Could not read any frames from that video.")
+    width, height = _pick_resolution(probe[0], preset)
 
-    width, height = _pick_resolution(frames[0], preset)
-    num_frames = int(num_frames)
+    ref = [_to_grayscale(f) for f in _load_frames(video, num_frames, width, height)]
+    prompt = _build_prompt(colors, scene, audio)
 
-    # Desaturate + resample + resize -> the grayscale reference the LoRA expects.
-    ref = [_to_grayscale(f).resize((width, height), Image.LANCZOS)
-           for f in _resample(frames, num_frames)]
+    def _cb(p, i, t, kw):
+        progress((i + 1) / NUM_STEPS, desc=f"Colorizing — step {i + 1}/{NUM_STEPS}")
+        return {}
 
-    prompt = _build_prompt(colors, scene)
-    ref_cond = LTX2ReferenceCondition(frames=ref, strength=1.0)
-
-    video_out, _audio = pipe(
+    video_out, audio_out = pipe(
         prompt=prompt,
         negative_prompt="",
-        reference_conditions=[ref_cond],
+        reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
         reference_downscale_factor=1,
-        width=width,
-        height=height,
-        num_frames=num_frames,
-        frame_rate=FPS,
-        num_inference_steps=NUM_STEPS,
-        sigmas=DISTILLED_SIGMA_VALUES,
-        guidance_scale=1.0,
-        stg_scale=0.0,
-        audio_guidance_scale=1.0,
-        audio_stg_scale=0.0,
+        width=width, height=height, num_frames=num_frames, frame_rate=FPS,
+        num_inference_steps=NUM_STEPS, sigmas=DISTILLED_SIGMA_VALUES,
+        guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
         generator=torch.Generator(device="cuda").manual_seed(seed),
-        output_type="np",
-        return_dict=False,
+        output_type="np", return_dict=False, callback_on_step_end=_cb,
     )
 
     out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
-    encode_video(video_out[0], fps=FPS, output_path=out_path)
+    _export(video_out[0], audio_out, out_path)
     return out_path, seed
 
 
@@ -142,25 +153,20 @@ def colorize(video, colors, scene, preset, num_frames, seed, randomize,
 with gr.Blocks(title="LTX-2.3 Colorize") as demo:
     gr.Markdown(
         "# 🎨 LTX-2.3 Video Colorization\n"
-        "Restore natural color to black-and-white or desaturated footage. "
-        "Upload a clip — it's converted to grayscale and recolored while keeping "
-        "subject, framing, and motion untouched. "
-        "IC-LoRA: [`linoyts/LTX-2.3-loras`](https://huggingface.co/linoyts/LTX-2.3-loras) · "
-        "base: distilled LTX-2.3."
+        "Restore natural color to black-and-white or desaturated footage. Upload a clip — it's converted "
+        "to grayscale and recolored while keeping subject, framing, and motion untouched. Optionally describe "
+        "the soundscape and the model generates matching audio. "
+        "IC-LoRA: [`linoyts/LTX-2.3-loras`](https://huggingface.co/linoyts/LTX-2.3-loras) · base: distilled LTX-2.3."
     )
     with gr.Row():
         with gr.Column():
             video_in = gr.Video(label="Input video (any clip — recolored as B&W)")
-            colors = gr.Textbox(
-                label="Colors to restore",
-                placeholder="a young brown rabbit with cream underside on grey granite rocks, warm afternoon light, green vegetation",
-                lines=2,
-            )
-            scene = gr.Textbox(
-                label="Scene description (optional)",
-                placeholder="a small rabbit sitting among boulders with dry grass behind it",
-                lines=2,
-            )
+            colors = gr.Textbox(label="Colors to restore", lines=2,
+                                placeholder="a young brown rabbit with cream underside on grey granite rocks, warm afternoon light, green vegetation")
+            scene = gr.Textbox(label="Scene description (optional)", lines=2,
+                               placeholder="a small rabbit sitting among boulders with dry grass behind it")
+            audio = gr.Textbox(label="Sound / audio (optional)", lines=1,
+                               placeholder="gentle wind, distant birdsong, soft rustling of dry grass")
             with gr.Accordion("Settings", open=False):
                 preset = gr.Dropdown(list(RES_PRESETS), value="Fast (768×448)", label="Resolution")
                 num_frames = gr.Dropdown(FRAME_CHOICES, value=73, label="Frames (24fps)")
@@ -171,38 +177,29 @@ with gr.Blocks(title="LTX-2.3 Colorize") as demo:
             video_out = gr.Video(label="Colorized result")
             used_seed = gr.Number(label="Seed used", interactive=False)
 
-    run.click(
-        colorize,
-        inputs=[video_in, colors, scene, preset, num_frames, seed, randomize],
-        outputs=[video_out, used_seed],
-    )
+    run.click(colorize, inputs=[video_in, colors, scene, audio, preset, num_frames, seed, randomize],
+              outputs=[video_out, used_seed])
 
     gr.Examples(
         examples=[
-            [
-                "examples/cat_on_a_tree_gray.mp4",
-                "a tabby cat with brown and grey fur and green eyes, clinging to rough brown tree bark, surrounded by lush green leaves",
-                "a cat climbing a tree among leaves",
-                "Fast (768×448)", 49, 42, False,
-            ],
-            [
-                "examples/surfing_gray.mp4",
-                "a surfer in a black wetsuit on a turquoise ocean wave, white foam, bright blue sky",
-                "a person surfing on an ocean wave",
-                "Fast (768×448)", 49, 42, False,
-            ],
-            [
-                "examples/slicing_veggie_gray.mp4",
-                "hands slicing a bright red tomato on a wooden cutting board, fresh green herbs and orange carrots nearby, warm kitchen light",
-                "hands slicing vegetables on a cutting board",
-                "Fast (768×448)", 49, 42, False,
-            ],
+            ["examples/slicing_veggie_gray.mp4",
+             "hands slicing a fresh green zucchini into thin rounds on a light wooden cutting board, the bright green skin and pale interior of the courgette, a stainless-steel knife catching warm kitchen light",
+             "close-up of hands slicing a vegetable on a cutting board",
+             "crisp rhythmic chopping on a wooden board, a faint kitchen ambience",
+             "Fast (768×448)", 73, 42, False],
+            ["examples/surfing_gray.mp4",
+             "a surfer in a black wetsuit riding a curling turquoise ocean wave, bright white foam spraying off the crest, deep blue sky and sunlit teal water",
+             "a person surfing along the face of an ocean wave",
+             "powerful ocean waves crashing and rushing water, wind, distant seagulls",
+             "Fast (768×448)", 73, 42, False],
+            ["examples/cat_on_a_tree_gray.mp4",
+             "a tabby cat with warm brown and grey striped fur and bright green eyes, clinging to rough red-brown tree bark surrounded by lush sunlit green leaves",
+             "a cat climbing a tree among leaves",
+             "soft leaves rustling in the breeze, faint birdsong, a quiet meow",
+             "Fast (768×448)", 73, 42, False],
         ],
-        inputs=[video_in, colors, scene, preset, num_frames, seed, randomize],
-        outputs=[video_out, used_seed],
-        fn=colorize,
-        cache_examples=True,
-        cache_mode="lazy",
+        inputs=[video_in, colors, scene, audio, preset, num_frames, seed, randomize],
+        outputs=[video_out, used_seed], fn=colorize, cache_examples=True, cache_mode="lazy",
     )
 
 if __name__ == "__main__":

examples/cat_on_a_tree_gray.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1cb0dda8f536f7e5e0beade829bed92cceebb8e605b0aaaff82a1f28a78a9f70
-size 214173
+oid sha256:15ca25a92785b9f942718e1f9411f668c605808b5c41dd933b5ae0be4e238224
+size 325396

examples/slicing_veggie_gray.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:20690e8a850cf08b8fa1e6e3038182421d72fed0a93d1900a372ba1a95bffa2a
-size 232034
+oid sha256:7707149cecf2ed4839904c8a5bf41105a4514dee88ac55f6f9761c537b822e9f
+size 351573

examples/surfing_gray.mp4

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ce4bf73761f71f382892d9c2ab058f716fd9b7b7e4a0c1107ed096242f52b551
-size 267897
+oid sha256:2b9b5b4b97d274d2dbfa5e829c2f7acc6510db4945b02f7bf7bb0e54200e0f39
+size 409218