import gc import os import subprocess import tempfile from dataclasses import dataclass from fractions import Fraction import av import gradio as gr import numpy as np import spaces import torch from PIL import Image from diffusers import LTX2InContextPipeline from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition from diffusers.pipelines.ltx2.utils import ( DEFAULT_NEGATIVE_PROMPT, DISTILLED_SIGMA_VALUES, ) MODEL_ID = "diffusers/LTX-2.3-Distilled-Diffusers" MODEL_REVISION = "432e0d3c2d1769aaa4d295f9243f7062bf6b47ee" LORA_ID = "Zlikwid/LTX_2.3_Upscale_IC_Lora" LORA_WEIGHT = "ltx2.3_upscale_ic-lora_06250.safetensors" REFERENCE_DOWNSCALE_FACTOR = 2 DISTILLED_STEPS = 8 DISTILLED_GUIDANCE_SCALE = 1.0 MAX_INPUT_DURATION_SECONDS = 12.0 MAX_UPLOAD_BYTES = 200 * 1024 * 1024 MAX_FRAME_COUNT = 65 MAX_FRAME_PIXELS = 1024 * 576 MAX_WORK_UNITS = MAX_FRAME_PIXELS * MAX_FRAME_COUNT FALLBACK_DECODE_FRAME_LIMIT = 900 INPUT_FACTOR_PRESET = "Source × Bicubic Factor" BICUBIC_RESAMPLE = ( Image.Resampling.BICUBIC if hasattr(Image, "Resampling") else Image.BICUBIC ) RESOLUTION_PRESETS = { INPUT_FACTOR_PRESET: (None, None), "768×512 (3:2)": (768, 512), "512×768 (2:3)": (512, 768), "1024×576 (16:9)": (1024, 576), "576×1024 (9:16)": (576, 1024), "768×768 (1:1)": (768, 768), "512×512 (1:1)": (512, 512), "960×544 (16:9)": (960, 544), "544×960 (9:16)": (544, 960), "Custom": (None, None), } PROCESSING_PRESETS = { "Preview": { "resolution": "512×512 (1:1)", "max_frames": 17, "lora_strength": 0.45, "input_preservation_strength": 0.90, }, "Balanced": { "resolution": "768×512 (3:2)", "max_frames": 41, "lora_strength": 0.55, "input_preservation_strength": 0.82, }, "Quality": { "resolution": "768×768 (1:1)", "max_frames": 65, "lora_strength": 0.65, "input_preservation_strength": 0.75, }, } print("Loading distilled LTX-2.3 pipeline...") pipe = LTX2InContextPipeline.from_pretrained( MODEL_ID, revision=MODEL_REVISION, torch_dtype=torch.bfloat16, ) print("Loading upscale IC-LoRA...") pipe.load_lora_weights( LORA_ID, adapter_name="upscale", weight_name=LORA_WEIGHT, ) pipe.vae.enable_tiling() print("Pipeline ready!") _offload_ready = False @dataclass class VideoMeta: path: str width: int height: int duration: float | None fps: float | None frame_count: int | None has_audio: bool file_size: int def video_path_from_input(input_video) -> str: if isinstance(input_video, dict): path = input_video.get("video") or input_video.get("path") elif isinstance(input_video, (list, tuple)): path = input_video[0] if input_video else None else: path = input_video if not path: raise gr.Error("Please upload a video first.") return str(path) def round_to_8k_plus_1(n: int) -> int: n = int(n) k = max(0, (n - 1) // 8) return max(9, k * 8 + 1) def snap_to_32(n: int) -> int: return max(32, round(n / 32) * 32) def format_duration(seconds: float | None) -> str: if seconds is None: return "unknown" return f"{seconds:.2f}s" def format_fps(fps: float | None) -> str: if fps is None: return "unknown" return f"{fps:.2f} fps" def update_stage_progress(progress, start, end, index, total, desc): if total <= 0: progress(end, desc=desc) return fraction = start + (end - start) * (index + 1) / total progress(fraction, desc=f"{desc} {index + 1}/{total}") def cleanup_cuda(): gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() def ensure_cpu_offload(): global _offload_ready if not _offload_ready: pipe.enable_model_cpu_offload() _offload_ready = True def probe_video(input_video) -> VideoMeta: path = video_path_from_input(input_video) if not os.path.exists(path): raise gr.Error(f"Input video does not exist: {path}") file_size = os.path.getsize(path) try: with av.open(path) as container: if not container.streams.video: raise gr.Error("Uploaded file has no video stream.") stream = container.streams.video[0] fps = float(stream.average_rate) if stream.average_rate else None duration = None if stream.duration is not None and stream.time_base is not None: duration = float(stream.duration * stream.time_base) elif container.duration is not None: duration = float(container.duration / 1_000_000) frame_count = int(stream.frames) if stream.frames else None if frame_count is None and duration and fps: frame_count = max(1, int(round(duration * fps))) width = int(stream.codec_context.width or stream.width) height = int(stream.codec_context.height or stream.height) has_audio = bool(container.streams.audio) except gr.Error: raise except Exception as e: raise gr.Error(f"Failed to inspect video: {e}") return VideoMeta( path=path, width=width, height=height, duration=duration, fps=fps, frame_count=frame_count, has_audio=has_audio, file_size=file_size, ) def resolve_output_size( meta: VideoMeta, resolution_preset: str, output_width, output_height, zoom, bicubic_upscale_factor, ) -> tuple[int, int, int, int]: zoom = float(zoom) bicubic_upscale_factor = float(bicubic_upscale_factor) source_width = max(32, int(meta.width / zoom)) if zoom > 1.0 else meta.width source_height = max(32, int(meta.height / zoom)) if zoom > 1.0 else meta.height if resolution_preset == INPUT_FACTOR_PRESET: width = snap_to_32(int(source_width * bicubic_upscale_factor)) height = snap_to_32(int(source_height * bicubic_upscale_factor)) elif resolution_preset == "Custom": width = snap_to_32(int(output_width)) height = snap_to_32(int(output_height)) else: width, height = RESOLUTION_PRESETS[resolution_preset] return width, height, source_width, source_height def validate_plan(meta: VideoMeta, output_width: int, output_height: int, num_frames: int): if meta.file_size > MAX_UPLOAD_BYTES: raise gr.Error( f"Input file is {meta.file_size / 1024 / 1024:.1f} MB. " f"Limit is {MAX_UPLOAD_BYTES / 1024 / 1024:.0f} MB." ) if meta.duration and meta.duration > MAX_INPUT_DURATION_SECONDS: raise gr.Error( f"Input duration is {meta.duration:.1f}s. " f"Limit is {MAX_INPUT_DURATION_SECONDS:.0f}s for this ZeroGPU Space." ) if meta.frame_count is not None and meta.frame_count < 9: raise gr.Error( f"Video has only {meta.frame_count} frames. Need at least 9 (8k+1)." ) if num_frames > MAX_FRAME_COUNT: raise gr.Error(f"Frame count {num_frames} exceeds limit {MAX_FRAME_COUNT}.") pixels = output_width * output_height if pixels > MAX_FRAME_PIXELS: raise gr.Error( f"Output {output_width}x{output_height} is too large for zero-a10g. " f"Limit is about {MAX_FRAME_PIXELS:,} pixels per frame." ) work_units = pixels * num_frames if work_units > MAX_WORK_UNITS: raise gr.Error( f"Requested {num_frames} frames at {output_width}x{output_height}. " "Reduce resolution or frame count." ) def planned_frame_count(meta: VideoMeta, max_frames) -> int: max_frames = min(MAX_FRAME_COUNT, int(max_frames)) available_frames = meta.frame_count or max_frames return round_to_8k_plus_1(min(max_frames, available_frames)) def output_fps_for_duration(meta: VideoMeta, num_frames: int) -> float: if meta.duration and meta.duration > 0: return max(0.1, (num_frames - 1) / meta.duration) if meta.fps: return meta.fps return 24.0 def encoded_fps_for_duration(meta: VideoMeta, frame_count: int) -> float: if meta.duration and meta.duration > 0: return max(0.1, frame_count / meta.duration) if meta.fps: return meta.fps return 24.0 def ffmpeg_rate_arg(fps: float) -> str: rate = Fraction(float(fps)).limit_denominator(100000) return f"{rate.numerator}/{rate.denominator}" def encoded_rate_arg_for_duration( meta: VideoMeta, frame_count: int, fallback_fps: float, ) -> str: if meta.duration and meta.duration > 0: duration = Fraction(float(meta.duration)).limit_denominator(10000) if duration > 0: rate = Fraction(frame_count, 1) / duration return f"{rate.numerator}/{rate.denominator}" return ffmpeg_rate_arg(fallback_fps) def normalize_output_frame(frame) -> np.ndarray: arr = np.asarray(frame) if arr.ndim != 3: raise gr.Error(f"Unexpected output frame shape: {arr.shape}") if arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): arr = np.moveaxis(arr, 0, -1) if arr.shape[-1] == 4: arr = arr[..., :3] elif arr.shape[-1] == 1: arr = np.repeat(arr, 3, axis=-1) elif arr.shape[-1] != 3: raise gr.Error(f"Unexpected output frame channels: {arr.shape}") if arr.dtype != np.uint8: arr = np.nan_to_num(arr) if arr.max(initial=0.0) <= 1.0 and arr.min(initial=0.0) >= 0.0: arr = arr * 255.0 arr = np.clip(arr, 0, 255).astype(np.uint8) return np.ascontiguousarray(arr) def encode_video_frames(frames, fps: float, rate_arg: str, output_path: str): frame_count = len(frames) if frame_count <= 0: raise gr.Error("Pipeline returned no output frames.") first_frame = normalize_output_frame(frames[0]) height, width = first_frame.shape[:2] raw_frames = [first_frame.tobytes()] for frame in frames[1:]: frame_arr = normalize_output_frame(frame) if frame_arr.shape[:2] != (height, width): raise gr.Error( f"Output frame size changed from {width}x{height} " f"to {frame_arr.shape[1]}x{frame_arr.shape[0]}." ) raw_frames.append(frame_arr.tobytes()) raw_video = b"".join(raw_frames) def command(use_demux_time_base: bool) -> list[str]: cmd = [ "ffmpeg", "-y", "-loglevel", "error", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s:v", f"{width}x{height}", "-framerate", rate_arg, "-i", "pipe:0", "-an", "-c:v", "libx264", "-preset", "medium", "-crf", "18", "-pix_fmt", "yuv420p", ] if use_demux_time_base: cmd.extend(["-enc_time_base", "demux"]) cmd.extend([ "-movflags", "+faststart", "-video_track_timescale", "90000", output_path, ]) return cmd errors = [] for use_demux_time_base in (True, False): result = subprocess.run( command(use_demux_time_base), input=raw_video, capture_output=True, ) if result.returncode == 0: return stderr = result.stderr.decode("utf-8", errors="replace") errors.append(stderr[-1000:]) raise gr.Error( "Failed to encode output video after ffmpeg retry: " + " | ".join(error for error in errors if error) ) def decode_sampled_frames( meta: VideoMeta, num_frames: int, progress=gr.Progress(), ) -> list[Image.Image]: progress(0.02, desc=f"Decoding {num_frames} sampled frames...") if meta.frame_count: indices = np.linspace(0, meta.frame_count - 1, num_frames).astype(int).tolist() wanted = set(indices) selected: dict[int, Image.Image] = {} with av.open(meta.path) as container: stream = container.streams.video[0] stream.thread_type = "AUTO" for frame_index, frame in enumerate(container.decode(stream)): if frame_index in wanted: selected[frame_index] = frame.to_image().convert("RGB") update_stage_progress( progress, 0.02, 0.08, len(selected) - 1, len(wanted), "Decoded sampled frames", ) if len(selected) == len(wanted): break frames = [selected[index] for index in indices if index in selected] if len(frames) == num_frames: return frames decoded = [] with av.open(meta.path) as container: stream = container.streams.video[0] stream.thread_type = "AUTO" for frame in container.decode(stream): decoded.append(frame.to_image().convert("RGB")) if len(decoded) > FALLBACK_DECODE_FRAME_LIMIT: raise gr.Error("Too many frames to decode safely in this Space.") if len(decoded) < 9: raise gr.Error(f"Video has only {len(decoded)} frames. Need at least 9 (8k+1).") num_frames = round_to_8k_plus_1(min(num_frames, len(decoded))) indices = np.linspace(0, len(decoded) - 1, num_frames).astype(int) frames = [decoded[index] for index in indices] progress(0.08, desc=f"Decoded {num_frames} sampled frames.") return frames def mux_source_audio(source_path: str, video_path: str, meta: VideoMeta, preserve_audio: bool): if not preserve_audio or not meta.has_audio: return video_path output_path = tempfile.mktemp(suffix=".mp4") cmd = [ "ffmpeg", "-y", "-i", video_path, "-i", source_path, "-map", "0:v:0", "-map", "1:a:0", "-c:v", "copy", "-c:a", "aac", "-shortest", output_path, ] result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: raise gr.Error(f"Failed to mux source audio: {result.stderr[-1000:]}") return output_path def estimate_settings( input_video, resolution_preset, output_width, output_height, zoom, bicubic_upscale_factor, max_frames, input_preservation_strength, preserve_audio, ): if input_video is None: return "Upload a video to see the planned output." try: meta = probe_video(input_video) num_frames = planned_frame_count(meta, max_frames) width, height, source_width, source_height = resolve_output_size( meta, resolution_preset, output_width, output_height, zoom, bicubic_upscale_factor, ) validate_plan(meta, width, height, num_frames) output_fps = output_fps_for_duration(meta, num_frames) except gr.Error as e: return f"Warning: {e}" except Exception as e: return f"Warning: failed to estimate settings: {e}" audio_status = "source audio" if preserve_audio and meta.has_audio else "silent/video only" return ( "**Run summary**\n\n" f"- Input: {meta.width}x{meta.height}, {format_duration(meta.duration)}, " f"{format_fps(meta.fps)}, {meta.frame_count or 'unknown'} frames\n" f"- Source after zoom: {source_width}x{source_height}\n" f"- Output: {width}x{height}, {num_frames} frames, " f"{output_fps:.2f} fps, duration preserved when metadata is available\n" f"- Model: distilled LTX 2.3, {DISTILLED_STEPS} steps, CFG {DISTILLED_GUIDANCE_SCALE}\n" f"- Input preservation: {float(input_preservation_strength):.2f}\n" f"- Audio: {audio_status}" ) @spaces.GPU(duration=600) def upscale_video( input_video, resolution_preset, output_width, output_height, zoom, bicubic_upscale_factor, lora_strength, input_preservation_strength, max_frames, seed, preserve_audio, progress=gr.Progress(track_tqdm=True), ): if input_video is None: raise gr.Error("Please upload a video first.") zoom = float(zoom) bicubic_upscale_factor = float(bicubic_upscale_factor) lora_strength = float(lora_strength) input_preservation_strength = float(input_preservation_strength) max_frames = int(max_frames) seed = int(seed) preserve_audio = bool(preserve_audio) progress(0.0, desc="Inspecting input video...") meta = probe_video(input_video) num_frames = planned_frame_count(meta, max_frames) output_width, output_height, _, _ = resolve_output_size( meta, resolution_preset, output_width, output_height, zoom, bicubic_upscale_factor, ) validate_plan(meta, output_width, output_height, num_frames) output_fps = output_fps_for_duration(meta, num_frames) try: frames = decode_sampled_frames(meta, num_frames, progress) progress(0.09, desc="Preparing bicubic pre-upscale...") if zoom > 1.0: crop_w = max(32, int(frames[0].width / zoom)) crop_h = max(32, int(frames[0].height / zoom)) left = (frames[0].width - crop_w) // 2 top = (frames[0].height - crop_h) // 2 cropped_frames = [] for index, frame in enumerate(frames): cropped_frames.append( frame.crop((left, top, left + crop_w, top + crop_h)) ) update_stage_progress( progress, 0.09, 0.12, index, len(frames), "Cropping frames" ) frames = cropped_frames else: progress(0.12, desc="No crop applied.") resized_frames = [] resize_desc = f"Bicubic pre-upscale to {output_width}x{output_height}" for index, frame in enumerate(frames): resized_frames.append( frame.resize((output_width, output_height), BICUBIC_RESAMPLE) ) update_stage_progress( progress, 0.12, 0.20, index, len(frames), resize_desc ) frames = resized_frames progress(0.21, desc="Building reference condition...") ref_cond = LTX2ReferenceCondition(frames=frames, strength=1.0) input_cond = LTX2VideoCondition( frames=frames, index=0, strength=input_preservation_strength, ) progress(0.23, desc="Applying LoRA strength...") pipe.set_adapters("upscale", lora_strength) progress(0.24, desc="Preparing GPU offload...") ensure_cpu_offload() generator = torch.Generator("cuda").manual_seed(seed) denoise_start = 0.26 denoise_end = 0.93 def step_callback(pipeline, step, timestep, kwargs): progress( denoise_start + (denoise_end - denoise_start) * (step + 1) / DISTILLED_STEPS, desc=f"Denoising step {step + 1}/{DISTILLED_STEPS}", ) return kwargs progress(denoise_start, desc=f"Denoising step 0/{DISTILLED_STEPS}") video, _audio = pipe( prompt="upscale", negative_prompt=DEFAULT_NEGATIVE_PROMPT, reference_conditions=[ref_cond], conditions=[input_cond], reference_downscale_factor=REFERENCE_DOWNSCALE_FACTOR, height=output_height, width=output_width, num_frames=num_frames, frame_rate=output_fps, num_inference_steps=DISTILLED_STEPS, sigmas=DISTILLED_SIGMA_VALUES, guidance_scale=DISTILLED_GUIDANCE_SCALE, stg_scale=1.0, modality_scale=3.0, guidance_rescale=0.7, spatio_temporal_guidance_blocks=[28], use_cross_timestep=True, generator=generator, output_type="np", return_dict=False, callback_on_step_end=step_callback, ) progress(0.95, desc="Encoding output video...") fd, video_path = tempfile.mkstemp(suffix=".mp4") os.close(fd) output_frames = video[0] actual_frame_count = len(output_frames) encoding_fps = encoded_fps_for_duration(meta, actual_frame_count) encoding_rate_arg = encoded_rate_arg_for_duration( meta, actual_frame_count, encoding_fps ) encode_video_frames(output_frames, encoding_fps, encoding_rate_arg, video_path) progress(0.97, desc="Preserving source audio..." if preserve_audio else "Finalizing...") output_path = mux_source_audio(meta.path, video_path, meta, preserve_audio) except torch.cuda.OutOfMemoryError: cleanup_cuda() raise gr.Error( "GPU out of memory. Reduce resolution or frame count. " "The ZeroGPU-safe limit is enforced, but this input still exceeded memory." ) except gr.Error: cleanup_cuda() raise except Exception as e: cleanup_cuda() raise gr.Error(f"Pipeline error: {e}") cleanup_cuda() progress( 1.0, desc=( f"Done: {output_width}x{output_height}, " f"{actual_frame_count} output frames, {encoding_fps:.2f} fps." ), ) return output_path with gr.Blocks( title="LTX 2.3 Distilled Upscale IC-LoRA", theme=gr.themes.Soft(primary_hue="blue", secondary_hue="purple"), ) as demo: gr.Markdown( """ # LTX 2.3 Distilled Video Upscaler Upload a short video to bicubic pre-upscale and refine it with the Zlikwid LTX 2.3 Upscale IC-LoRA. """ ) with gr.Row(): with gr.Column(scale=1): input_video = gr.Video(label="Input Video", sources=["upload"]) run_summary = gr.Markdown("Upload a video to see the planned output.") with gr.Accordion("Settings", open=True): processing_preset = gr.Radio( choices=list(PROCESSING_PRESETS.keys()), value="Balanced", label="Processing Preset", ) resolution_preset = gr.Dropdown( choices=list(RESOLUTION_PRESETS.keys()), value="768×512 (3:2)", label="Output Resolution", ) with gr.Row(): output_width = gr.Slider( 64, 1024, value=768, step=32, label="Width (px)", visible=False, ) output_height = gr.Slider( 64, 1024, value=512, step=32, label="Height (px)", visible=False, ) zoom = gr.Slider( 1.0, 4.0, value=1.0, step=0.1, label="Zoom", info="1.0 keeps the full frame; higher values crop the center.", ) bicubic_upscale_factor = gr.Slider( 1.0, 4.0, value=2.0, step=0.25, label="Bicubic Pre-Upscale Factor", info="Used when Output Resolution is Source × Bicubic Factor.", visible=False, ) lora_strength = gr.Slider( 0.0, 1.0, value=0.55, step=0.05, label="Refinement Strength (LoRA)", info="Higher values add more generative detail; lower values stay closer to input.", ) input_preservation_strength = gr.Slider( 0.0, 1.0, value=0.82, step=0.05, label="Input Preservation Strength", info="Higher values anchor the output to the bicubic input more strongly.", ) max_frames = gr.Slider( 9, MAX_FRAME_COUNT, value=41, step=8, label="Sampled Frames (8k+1)", ) with gr.Row(): seed = gr.Number( value=42, label="Seed", precision=0, ) randomize_seed = gr.Checkbox( label="Randomize", value=False, ) preserve_audio = gr.Checkbox( label="Preserve Source Audio", value=True, ) generate_btn = gr.Button( "Upscale Video", variant="primary", size="lg", ) with gr.Column(scale=1): output_video = gr.Video(label="Output Video") def controls_for_resolution(preset): if preset == "Custom": return ( gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), ) if preset == INPUT_FACTOR_PRESET: return ( gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), ) return ( gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), ) def toggle_custom_resolution(preset): return controls_for_resolution(preset) resolution_preset.change( toggle_custom_resolution, inputs=[resolution_preset], outputs=[output_width, output_height, bicubic_upscale_factor], queue=False, show_progress="hidden", ) def apply_processing_preset( preset, input_video_value, output_width_value, output_height_value, zoom_value, bicubic_upscale_factor_value, preserve_audio_value, ): values = PROCESSING_PRESETS[preset] resolution = values["resolution"] width_update, height_update, factor_update = controls_for_resolution(resolution) summary = estimate_settings( input_video_value, resolution, output_width_value, output_height_value, zoom_value, bicubic_upscale_factor_value, values["max_frames"], values["input_preservation_strength"], preserve_audio_value, ) return ( gr.update(value=resolution), width_update, height_update, factor_update, gr.update(value=values["max_frames"]), gr.update(value=values["lora_strength"]), gr.update(value=values["input_preservation_strength"]), summary, ) processing_preset.change( apply_processing_preset, inputs=[ processing_preset, input_video, output_width, output_height, zoom, bicubic_upscale_factor, preserve_audio, ], outputs=[ resolution_preset, output_width, output_height, bicubic_upscale_factor, max_frames, lora_strength, input_preservation_strength, run_summary, ], queue=False, show_progress="hidden", ) summary_inputs = [ input_video, resolution_preset, output_width, output_height, zoom, bicubic_upscale_factor, max_frames, input_preservation_strength, preserve_audio, ] for component in [ input_video, resolution_preset, zoom, bicubic_upscale_factor, max_frames, input_preservation_strength, preserve_audio, output_width, output_height, ]: component.change( estimate_settings, inputs=summary_inputs, outputs=[run_summary], queue=False, show_progress="hidden", ) def maybe_randomize(seed_val, do_randomize): if do_randomize: return int(np.random.randint(0, 2**32)) return seed_val generate_btn.click( maybe_randomize, inputs=[seed, randomize_seed], outputs=[seed], queue=False, show_progress="hidden", ).then( upscale_video, inputs=[ input_video, resolution_preset, output_width, output_height, zoom, bicubic_upscale_factor, lora_strength, input_preservation_strength, max_frames, seed, preserve_audio, ], outputs=[output_video], concurrency_limit=1, ) demo.queue(default_concurrency_limit=1, max_size=4) if __name__ == "__main__": demo.launch()