Spaces:
Running on Zero
Running on Zero
multimodalart HF Staff
Fix: pass drag-widget JS via gr.Blocks(head=...) not launch()
47929ce verified | """LTX-2.3 Relight (Video Relighting IC-LoRA) demo. | |
| Relights an exterior video clip using the Lightricks LTX-2.3-22B model with the | |
| ``Lightricks/LTX-2.3-22b-IC-LoRA-Relight`` adapter. The user picks a light | |
| direction (azimuth / elevation / hardness) and a lighting "look"; the demo | |
| renders a light-direction sphere, composites it into the top-right corner of | |
| every frame, and feeds the result as the IC-LoRA reference conditioning stream. | |
| This is a faithful port of the author's ComfyUI single-stage distilled V2V | |
| IC-LoRA workflow | |
| (``LTX-2.3_Relight_ICLoRA_SingleStage_Distilled.json``) onto the official | |
| ``ltx_pipelines`` / ``ltx_core`` packages (vendored under ./packages). | |
| """ | |
| import os | |
| import sys | |
| from pathlib import Path | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") | |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") | |
| # ZeroGPU: torch.compile / dynamo unsupported. | |
| os.environ["TORCH_COMPILE_DISABLE"] = "1" | |
| os.environ["TORCHDYNAMO_DISABLE"] = "1" | |
| # Vendored LTX-2 packages on sys.path *before* any ltx_core / ltx_pipelines import. | |
| _HERE = Path(__file__).parent | |
| sys.path.insert(0, str(_HERE / "packages" / "ltx-pipelines" / "src")) | |
| sys.path.insert(0, str(_HERE / "packages" / "ltx-core" / "src")) | |
| import base64 | |
| import io | |
| import json | |
| import logging | |
| import tempfile | |
| from fractions import Fraction | |
| import spaces | |
| import gradio as gr | |
| import torch | |
| torch._dynamo.config.suppress_errors = True | |
| torch._dynamo.config.disable = True | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| # --- LTX-2 imports (vendored copy under ./packages) -------------------------- | |
| from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number | |
| from ltx_core.quantization.fp8_cast import build_policy as build_fp8_cast_policy | |
| from ltx_core.loader import LTXV_LORA_COMFY_RENAMING_MAP, LoraPathStrengthAndSDOps | |
| from ltx_pipelines.ic_lora import ICLoraPipeline | |
| from ltx_pipelines.utils.constants import DISTILLED_SIGMAS | |
| from ltx_pipelines.utils.media_io import encode_video | |
| # --- attention backend patch (match the compile Space's block build) --------- | |
| # SDPA instead of FA3 (FA3 crashes on Blackwell ZeroGPU; SDPA is portable). | |
| import torch.nn.functional as _F | |
| from ltx_core.model.transformer import attention as _attn_mod | |
| def _sdpa_as_mea(query, key, value, attn_bias=None, scale=None, **kwargs): | |
| q, k, v = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) | |
| return _F.scaled_dot_product_attention(q, k, v, scale=scale).transpose(1, 2) | |
| _attn_mod.memory_efficient_attention = _sdpa_as_mea | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger("relight") | |
| # ----------------------------------------------------------------------------- | |
| # Constants (match the author's ComfyUI V2V IC-LoRA single-stage distilled workflow) | |
| # ----------------------------------------------------------------------------- | |
| BASE_REPO = "Lightricks/LTX-2.3" | |
| DISTILLED_CHECKPOINT_FILE = "ltx-2.3-22b-distilled.safetensors" | |
| SPATIAL_UPSAMPLER_FILE = "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" | |
| LORA_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-Relight" | |
| LORA_FILE = "ltx-2.3-22b-ic-lora-relight-1.0.safetensors" | |
| # Distilled LoRA (used alongside the IC-LoRA at strength 0.5 in the workflow) | |
| DISTILLED_LORA_FILE = "ltx-2.3-22b-distilled-lora-384-1.1.safetensors" | |
| GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized" | |
| TARGET_FPS = 24 | |
| FRAME_MULTIPLE = 8 | |
| # Cap frames to keep a single ZeroGPU allocation bounded (~5s @ 24fps = 121 frames). | |
| MAX_FRAMES = 121 | |
| MIN_DURATION = 0.5 | |
| MAX_DURATION = round((MAX_FRAMES - 1) / TARGET_FPS, 2) # 5.0 s | |
| DEFAULT_DURATION = 3.0 | |
| # The relight LoRA was trained at 1280x704. The pipeline trick from the model | |
| # card: request 2560x1408 and skip stage 2 to output 1280x704. We keep a single | |
| # fixed resolution because the sphere overlay geometry is calibrated for 1280x704. | |
| TARGET_WIDTH = 1280 | |
| TARGET_HEIGHT = 704 | |
| # Sphere overlay geometry (at 1280x704, from the model card / ComfyUI workflow) | |
| SPHERE_PATCH_SIZE = 143 | |
| SPHERE_MARGIN = 22 # px from top and right | |
| TRIGGER_PHRASE = "relight the video to match the light-direction ball." | |
| # The 12 trained lighting "looks" from the model card. | |
| LIGHT_LOOKS = [ | |
| "hard directional sunlight", | |
| "hard high-angle sunlight", | |
| "hard low-angle sunlight", | |
| "soft diffused daylight", | |
| "soft warm afternoon light", | |
| "cool soft daylight", | |
| "dim overcast light", | |
| "strong backlight with rim light", | |
| "soft hazy backlight", | |
| "warm golden low front sun", | |
| "warm golden low side sun", | |
| "frontal sunlight", | |
| ] | |
| DEFAULT_LIGHT_LOOK = "hard directional sunlight" | |
| # Default light direction (matches the workflow's SphereLightNode: az≈150, el≈15) | |
| DEFAULT_AZIMUTH = 150.0 | |
| DEFAULT_ELEVATION = 15.0 | |
| DEFAULT_HARDNESS = 2 | |
| DEFAULT_NEGATIVE = ( | |
| "pc game, console game, video game, cartoon, childish, ugly, blurry, " | |
| "deformed, inconsistent motion, worst quality, jittery" | |
| ) | |
| DTYPE = torch.bfloat16 | |
| DEVICE = torch.device("cuda") | |
| # ----------------------------------------------------------------------------- | |
| # Sphere-light renderer (pure Python, matching the ComfyUI SphereLightNode) | |
| # ----------------------------------------------------------------------------- | |
| # Replicates the Three.js scene from the SphereLightNode custom node: | |
| # - Camera at (0, 6, 8) looking at (0, -0.5, 0), FOV 35° | |
| # - Unit sphere at origin, gray material (RGB 204), roughness 0.8 | |
| # - Ground plane at y=-1, gray (RGB 138), receives shadows | |
| # - Directional light at position computed from azimuth/elevation | |
| # - Ambient light intensity 0.2 | |
| # | |
| # The `hardness` widget from the ComfyUI workflow maps to light intensity: | |
| # 1 = soft (intensity ~0.8), 2 = medium (intensity ~1.5), 3 = hard (intensity ~2.5) | |
| import numpy as np | |
| from PIL import Image, ImageFilter | |
| _RENDER_SIZE = 512 | |
| _BG_RGB = np.array([138, 138, 138], dtype=np.float32) / 255.0 | |
| _SPHERE_RGB = np.array([204, 204, 204], dtype=np.float32) / 255.0 | |
| _HARDNESS_TO_INTENSITY = {1: 0.8, 2: 1.5, 3: 2.5} | |
| def _normalize(vec): | |
| vec = np.asarray(vec, dtype=np.float32) | |
| norm = np.linalg.norm(vec, axis=-1, keepdims=True) | |
| return vec / np.maximum(norm, 1e-6) | |
| def _build_camera_rays(size): | |
| cam_pos = np.array([0.0, 6.0, 8.0], dtype=np.float32) | |
| cam_target = np.array([0.0, -0.5, 0.0], dtype=np.float32) | |
| world_up = np.array([0.0, 1.0, 0.0], dtype=np.float32) | |
| forward = _normalize(cam_target - cam_pos) | |
| right = _normalize(np.cross(forward, world_up)) | |
| up = _normalize(np.cross(right, forward)) | |
| yy, xx = np.mgrid[0:size, 0:size].astype(np.float32) | |
| nx = ((xx + 0.5) / size) * 2.0 - 1.0 | |
| ny = 1.0 - ((yy + 0.5) / size) * 2.0 | |
| scale = np.tan(np.deg2rad(np.float32(35.0)) / 2.0) | |
| dirs = ( | |
| forward[None, None, :] | |
| + nx[..., None] * scale * right[None, None, :] | |
| + ny[..., None] * scale * up[None, None, :] | |
| ) | |
| return cam_pos, _normalize(dirs) | |
| def _intersect_sphere(ray_origin, ray_dirs, center, radius): | |
| oc = ray_origin[None, None, :] - center[None, None, :] | |
| half_b = np.sum(ray_dirs * oc, axis=-1) | |
| c = np.sum(oc * oc, axis=-1) - radius * radius | |
| disc = half_b * half_b - c | |
| valid = disc > 0.0 | |
| sqrt_disc = np.zeros_like(disc, dtype=np.float32) | |
| sqrt_disc[valid] = np.sqrt(disc[valid]) | |
| t = np.full(disc.shape, np.inf, dtype=np.float32) | |
| near = -half_b - sqrt_disc | |
| far = -half_b + sqrt_disc | |
| valid_near = valid & (near > 1e-4) | |
| valid_far = valid & ~valid_near & (far > 1e-4) | |
| t[valid_near] = near[valid_near] | |
| t[valid_far] = far[valid_far] | |
| return t | |
| def _intersect_plane_y(ray_origin, ray_dirs, plane_y): | |
| dy = ray_dirs[..., 1] | |
| t = np.full(dy.shape, np.inf, dtype=np.float32) | |
| valid = np.abs(dy) > 1e-6 | |
| raw_t = (plane_y - ray_origin[1]) / dy | |
| t[valid & (raw_t > 1e-4)] = raw_t[valid & (raw_t > 1e-4)] | |
| return t | |
| def _batched_directional_shadow(points, light_dir, sphere_radius): | |
| if points.size == 0: | |
| return np.zeros((0,), dtype=np.float32) | |
| helper = np.array([0.0, 1.0, 0.0], dtype=np.float32) | |
| if abs(np.dot(helper, light_dir)) > 0.95: | |
| helper = np.array([1.0, 0.0, 0.0], dtype=np.float32) | |
| tangent_a = _normalize(np.cross(light_dir, helper)).reshape(3) | |
| tangent_b = _normalize(np.cross(light_dir, tangent_a)).reshape(3) | |
| offsets = [(0, 0), (0.055, 0), (-0.055, 0), (0, 0.055), (0, -0.055)] | |
| occlusion = np.zeros((points.shape[0],), dtype=np.float32) | |
| for off_a, off_b in offsets: | |
| sample_dir = _normalize(light_dir + tangent_a * off_a + tangent_b * off_b).reshape(3) | |
| origin = points + sample_dir[None, :] * 0.03 | |
| proj = np.sum(origin * sample_dir[None, :], axis=1) | |
| c = np.sum(origin * origin, axis=1) - sphere_radius * sphere_radius | |
| disc = proj * proj - c | |
| hit = disc > 0.0 | |
| sqrt_disc = np.zeros_like(disc, dtype=np.float32) | |
| sqrt_disc[hit] = np.sqrt(disc[hit]) | |
| t = -proj - sqrt_disc | |
| occlusion += (hit & (t > 1e-4)).astype(np.float32) | |
| return occlusion / float(len(offsets)) | |
| def render_sphere_image(azimuth, elevation, hardness): | |
| """Render a light-direction sphere image (PIL Image, 1024x1024 RGB). | |
| Args: | |
| azimuth: Light azimuth in degrees (0 = front, 90 = right, 180 = behind). | |
| elevation: Light elevation in degrees (0 = horizon, 90 = top). | |
| hardness: 1=soft, 2=medium, 3=hard (maps to directional light intensity). | |
| """ | |
| intensity = _HARDNESS_TO_INTENSITY.get(int(hardness), 1.5) | |
| cam_pos, ray_dirs = _build_camera_rays(_RENDER_SIZE) | |
| sphere_center = np.array([0.0, 0.0, 0.0], dtype=np.float32) | |
| sphere_radius = 1.0 | |
| plane_y = -1.0 | |
| az = np.deg2rad(np.float32(azimuth)) | |
| el = np.deg2rad(np.float32(elevation)) | |
| light_pos = np.array( | |
| [ | |
| 10.0 * np.cos(el) * np.sin(az), | |
| 10.0 * np.sin(el), | |
| 10.0 * np.cos(el) * np.cos(az), | |
| ], | |
| dtype=np.float32, | |
| ) | |
| light_dir = _normalize(light_pos).reshape(3) | |
| sphere_t = _intersect_sphere(cam_pos, ray_dirs, sphere_center, sphere_radius) | |
| plane_t = _intersect_plane_y(cam_pos, ray_dirs, plane_y) | |
| image = np.broadcast_to(_BG_RGB, (_RENDER_SIZE, _RENDER_SIZE, 3)).copy() | |
| sphere_hit = np.isfinite(sphere_t) & (sphere_t <= plane_t) | |
| plane_hit = np.isfinite(plane_t) & ~sphere_hit | |
| if np.any(sphere_hit): | |
| sphere_points = cam_pos[None, None, :] + ray_dirs * sphere_t[..., None] | |
| sphere_points_hit = sphere_points[sphere_hit] | |
| normals_hit = _normalize(sphere_points_hit - sphere_center[None, :]) | |
| view_dir_hit = _normalize(cam_pos[None, :] - sphere_points_hit) | |
| diffuse = np.clip(np.sum(normals_hit * light_dir[None, :], axis=1), 0.0, 1.0) | |
| half_vec = _normalize(light_dir[None, :] + view_dir_hit) | |
| specular = np.clip(np.sum(normals_hit * half_vec, axis=1), 0.0, 1.0) ** 28.0 | |
| ambient = 0.2 | |
| diffuse_term = 0.68 * diffuse * float(intensity) | |
| specular_term = 0.14 * specular * min(float(intensity), 2.0) | |
| sphere_color = _SPHERE_RGB[None, :] * np.clip( | |
| (ambient + diffuse_term)[:, None], 0.0, 1.35 | |
| ) | |
| sphere_color = np.clip(sphere_color + specular_term[:, None], 0.0, 1.0) | |
| image[sphere_hit] = sphere_color | |
| if np.any(plane_hit): | |
| plane_points = cam_pos[None, None, :] + ray_dirs * plane_t[..., None] | |
| plane_points_hit = plane_points[plane_hit] | |
| occlusion = _batched_directional_shadow(plane_points_hit, light_dir, sphere_radius) | |
| shadow_strength = 0.42 + 0.18 * min(float(intensity) / 3.0, 1.0) | |
| plane_color = _BG_RGB[None, :] * (1.0 - shadow_strength * occlusion[:, None]) | |
| image[plane_hit] = np.clip(plane_color, 0.0, 1.0) | |
| rgb = np.clip(image * 255.0, 0.0, 255.0).astype(np.uint8) | |
| img = Image.fromarray(rgb, mode="RGB").resize((1024, 1024), Image.LANCZOS) | |
| img = img.filter(ImageFilter.GaussianBlur(radius=1.5)) | |
| img = img.point(lambda value: min(255, value + 40)) | |
| return img | |
| # ----------------------------------------------------------------------------- | |
| # Video preprocessing: composite the light-direction sphere into every frame | |
| # ----------------------------------------------------------------------------- | |
| def _probe_video(video_path): | |
| """Return (num_frames_at_source_fps, src_fps, width, height).""" | |
| import av | |
| container = av.open(video_path) | |
| try: | |
| stream = next(s for s in container.streams if s.type == "video") | |
| w = stream.codec_context.width | |
| h = stream.codec_context.height | |
| fps = float(stream.average_rate) if stream.average_rate else TARGET_FPS | |
| frames = stream.frames | |
| if not frames or frames <= 0: | |
| frames = sum(1 for _ in container.decode(stream)) | |
| return frames, fps, w, h | |
| finally: | |
| container.close() | |
| def _plan_num_frames(video_path, duration): | |
| """Compute the output frame count aligned to (1 + 8k) @ 24fps.""" | |
| frames, fps, w, h = _probe_video(video_path) | |
| if fps and fps > 0: | |
| avail = int(round(frames * TARGET_FPS / fps)) | |
| else: | |
| avail = frames | |
| n = min(avail if avail > 0 else MAX_FRAMES, MAX_FRAMES) | |
| if duration is not None and duration > 0: | |
| requested = int(round(float(duration) * TARGET_FPS)) | |
| n = min(n, max(requested, 1)) | |
| n = ((n - 1) // FRAME_MULTIPLE) * FRAME_MULTIPLE + 1 | |
| n = max(n, 1 + FRAME_MULTIPLE) | |
| return n | |
| def _decode_ball_snapshot(ball_state): | |
| """Decode the interactive widget's rendered sphere PNG (data-URL) into a | |
| 143x143 RGB PIL patch, ready to composite. | |
| ``ball_state`` is the JSON string written by the drag-to-set light-direction | |
| widget: ``{"azimuth": .., "elevation": .., "hardness": .., "png": "data:image/png;base64,.."}``. | |
| Returns None if no valid PNG snapshot is present (caller should fall back to | |
| the pure-Python render). | |
| """ | |
| if not ball_state: | |
| return None | |
| try: | |
| state = json.loads(ball_state) | |
| except (TypeError, ValueError): | |
| return None | |
| png = state.get("png", "") if isinstance(state, dict) else "" | |
| if not isinstance(png, str) or not png.startswith("data:image"): | |
| return None | |
| try: | |
| raw = base64.b64decode(png.split(",", 1)[1]) | |
| ball = Image.open(io.BytesIO(raw)).convert("RGB") | |
| except Exception as exc: # noqa: BLE001 | |
| logger.warning("Failed to decode ball snapshot: %s", exc) | |
| return None | |
| return ball.resize((SPHERE_PATCH_SIZE, SPHERE_PATCH_SIZE), Image.LANCZOS) | |
| def _build_sphere_overlay(azimuth, elevation, hardness, ball_state=None): | |
| """Return the 143x143 RGB sphere patch to composite. | |
| Prefers the snapshot rendered by the interactive drag-to-set widget | |
| (``ball_state``); falls back to the pure-Python ray-traced render when the | |
| widget hasn't produced one (e.g. headless gradio_client / API calls). | |
| """ | |
| patch = _decode_ball_snapshot(ball_state) | |
| if patch is not None: | |
| return patch | |
| sphere_img = render_sphere_image(azimuth, elevation, hardness) | |
| # Resize to the exact patch size used by the workflow | |
| sphere_img = sphere_img.resize((SPHERE_PATCH_SIZE, SPHERE_PATCH_SIZE), Image.LANCZOS) | |
| return sphere_img | |
| def _composite_sphere_into_frames(video_path, num_frames, azimuth, elevation, | |
| hardness, ball_state=None): | |
| """Re-encode the source to 24fps at 1280x704 with the light-direction ball | |
| composited into the top-right corner of every frame. | |
| Returns the path to the processed reference video. | |
| """ | |
| import av | |
| sphere_patch = _build_sphere_overlay(azimuth, elevation, hardness, ball_state) | |
| # Convert to numpy for fast compositing | |
| sphere_arr = np.array(sphere_patch, dtype=np.uint8) # (143, 143, 3) | |
| in_container = av.open(video_path) | |
| resampled = [] | |
| try: | |
| in_stream = next(s for s in in_container.streams if s.type == "video") | |
| src_w = in_stream.codec_context.width | |
| src_h = in_stream.codec_context.height | |
| src_fps = float(in_stream.average_rate) if in_stream.average_rate else TARGET_FPS | |
| # Scale to 1280x704 (preserve aspect, center crop/pad) | |
| # The workflow uses: ResizeImageMaskNode (scale shorter dim to 704) then | |
| # ImageScale (to 1280x704, center). We replicate: scale so shorter dim | |
| # becomes 704, then center-crop/pad to 1280x704. | |
| if src_h < src_w: | |
| # Landscape: scale height to 704, width scales proportionally | |
| scale_h = TARGET_HEIGHT | |
| scale_w = int(src_w * scale_h / src_h) | |
| else: | |
| # Portrait/square: scale width to ... actually scale shorter dim (width) to 704? | |
| # No — the workflow scales "shorter dimension" to 704. For portrait, shorter=width. | |
| # But the output must be 1280x704 (landscape). So we scale shorter dim to 704 | |
| # and then resize to 1280x704. For landscape, shorter=height→704. | |
| scale_w = TARGET_WIDTH | |
| scale_h = int(src_h * scale_w / src_w) | |
| step = 1.0 / TARGET_FPS | |
| next_t = 0.0 | |
| idx = 0 | |
| for frame in in_container.decode(in_stream): | |
| t = float(frame.pts * in_stream.time_base) if frame.pts is not None else idx / src_fps | |
| idx += 1 | |
| while next_t <= t and len(resampled) < num_frames: | |
| # Reformat to target size | |
| reformatted = frame.reformat(width=scale_w, height=scale_h, format="rgb24") | |
| arr = np.array(reformatted.to_ndarray(), dtype=np.uint8) if hasattr(reformatted, 'to_ndarray') else np.frombuffer(reformatted.planes[0], dtype=np.uint8).reshape(scale_h, scale_w, 3) | |
| # Center crop/pad to 1280x704 | |
| if arr.shape[1] != TARGET_WIDTH or arr.shape[0] != TARGET_HEIGHT: | |
| canvas = np.full((TARGET_HEIGHT, TARGET_WIDTH, 3), 128, dtype=np.uint8) | |
| y_off = max(0, (TARGET_HEIGHT - arr.shape[0]) // 2) | |
| x_off = max(0, (TARGET_WIDTH - arr.shape[1]) // 2) | |
| h_copy = min(arr.shape[0], TARGET_HEIGHT) | |
| w_copy = min(arr.shape[1], TARGET_WIDTH) | |
| canvas[max(0, y_off):max(0, y_off) + h_copy, | |
| max(0, x_off):max(0, x_off) + w_copy] = arr[:h_copy, :w_copy] | |
| # If source is larger, crop from center | |
| if arr.shape[0] > TARGET_HEIGHT or arr.shape[1] > TARGET_WIDTH: | |
| cy = (arr.shape[0] - TARGET_HEIGHT) // 2 | |
| cx = (arr.shape[1] - TARGET_WIDTH) // 2 | |
| canvas = arr[cy:cy + TARGET_HEIGHT, cx:cx + TARGET_WIDTH].copy() | |
| arr = canvas | |
| # Composite the sphere into the top-right corner | |
| # Position: 22px from top, 22px from right | |
| x_pos = TARGET_WIDTH - SPHERE_PATCH_SIZE - SPHERE_MARGIN # 1280-143-22 = 1115 | |
| y_pos = SPHERE_MARGIN # 22 | |
| arr[y_pos:y_pos + SPHERE_PATCH_SIZE, x_pos:x_pos + SPHERE_PATCH_SIZE] = sphere_arr | |
| # Convert back to av.VideoFrame | |
| out_frame = av.VideoFrame.from_ndarray(arr, format="rgb24") | |
| resampled.append(out_frame) | |
| next_t += step | |
| if len(resampled) >= num_frames: | |
| break | |
| while resampled and len(resampled) < num_frames: | |
| resampled.append(resampled[-1]) | |
| finally: | |
| in_container.close() | |
| if not resampled: | |
| raise gr.Error("Could not decode any frames from the input video.") | |
| out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name | |
| out_container = av.open(out_path, mode="w") | |
| try: | |
| vstream = out_container.add_stream("libx264", rate=TARGET_FPS) | |
| vstream.width = TARGET_WIDTH | |
| vstream.height = TARGET_HEIGHT | |
| vstream.pix_fmt = "yuv420p" | |
| vstream.time_base = Fraction(1, TARGET_FPS) | |
| vstream.options = {"crf": "18"} | |
| # Silent stereo audio track (the LTX-2 loader expects audio) | |
| astream = out_container.add_stream("aac", rate=48000) | |
| astream.layout = "stereo" | |
| for i, out_frame in enumerate(resampled): | |
| out_frame.pts = i | |
| out_frame.time_base = Fraction(1, TARGET_FPS) | |
| for pkt in vstream.encode(out_frame): | |
| out_container.mux(pkt) | |
| for pkt in vstream.encode(): | |
| out_container.mux(pkt) | |
| # Encode silence for the full duration | |
| duration = len(resampled) / TARGET_FPS | |
| samples_per_frame = 1024 | |
| total_samples = int(48000 * duration) | |
| pts = 0 | |
| while pts < total_samples: | |
| n = min(samples_per_frame, total_samples - pts) | |
| arr = np.zeros((2, n), dtype=np.float32) | |
| aframe = av.AudioFrame.from_ndarray(arr, format="fltp", layout="stereo") | |
| aframe.sample_rate = 48000 | |
| aframe.pts = pts | |
| aframe.time_base = Fraction(1, 48000) | |
| for pkt in astream.encode(aframe): | |
| out_container.mux(pkt) | |
| pts += n | |
| for pkt in astream.encode(): | |
| out_container.mux(pkt) | |
| finally: | |
| out_container.close() | |
| return out_path | |
| def _preview_composite(video_path, azimuth, elevation, hardness, ball_state=None): | |
| """Render a single preview frame showing the sphere composited into the | |
| first frame of the video. Returns a PIL Image (for gr.Image preview).""" | |
| import av | |
| sphere_patch = _build_sphere_overlay(azimuth, elevation, hardness, ball_state) | |
| sphere_arr = np.array(sphere_patch, dtype=np.uint8) | |
| container = av.open(video_path) | |
| try: | |
| stream = next(s for s in container.streams if s.type == "video") | |
| src_w = stream.codec_context.width | |
| src_h = stream.codec_context.height | |
| for frame in container.decode(stream): | |
| # Scale to 1280x704 | |
| if src_h < src_w: | |
| scale_h = TARGET_HEIGHT | |
| scale_w = int(src_w * scale_h / src_h) | |
| else: | |
| scale_w = TARGET_WIDTH | |
| scale_h = int(src_h * scale_w / src_w) | |
| reformatted = frame.reformat(width=scale_w, height=scale_h, format="rgb24") | |
| arr = np.array(reformatted.to_ndarray(), dtype=np.uint8) if hasattr(reformatted, 'to_ndarray') else np.frombuffer(reformatted.planes[0], dtype=np.uint8).reshape(scale_h, scale_w, 3) | |
| # Center crop/pad to 1280x704 | |
| if arr.shape[1] != TARGET_WIDTH or arr.shape[0] != TARGET_HEIGHT: | |
| if arr.shape[0] > TARGET_HEIGHT or arr.shape[1] > TARGET_WIDTH: | |
| cy = (arr.shape[0] - TARGET_HEIGHT) // 2 | |
| cx = (arr.shape[1] - TARGET_WIDTH) // 2 | |
| arr = arr[cy:cy + TARGET_HEIGHT, cx:cx + TARGET_WIDTH].copy() | |
| else: | |
| canvas = np.full((TARGET_HEIGHT, TARGET_WIDTH, 3), 128, dtype=np.uint8) | |
| y_off = (TARGET_HEIGHT - arr.shape[0]) // 2 | |
| x_off = (TARGET_WIDTH - arr.shape[1]) // 2 | |
| canvas[y_off:y_off + arr.shape[0], x_off:x_off + arr.shape[1]] = arr | |
| arr = canvas | |
| # Composite sphere | |
| x_pos = TARGET_WIDTH - SPHERE_PATCH_SIZE - SPHERE_MARGIN | |
| y_pos = SPHERE_MARGIN | |
| arr[y_pos:y_pos + SPHERE_PATCH_SIZE, x_pos:x_pos + SPHERE_PATCH_SIZE] = sphere_arr | |
| return Image.fromarray(arr, mode="RGB") | |
| finally: | |
| container.close() | |
| return None | |
| # ----------------------------------------------------------------------------- | |
| # Download weights at startup (cached on the Space) | |
| # ----------------------------------------------------------------------------- | |
| _HF_TOKEN = os.environ.get("HF_TOKEN") | |
| logger.info("Downloading LTX-2.3 distilled checkpoint (~46GB) ...") | |
| CHECKPOINT_PATH = hf_hub_download(BASE_REPO, DISTILLED_CHECKPOINT_FILE, token=_HF_TOKEN) | |
| logger.info("Downloading spatial upsampler ...") | |
| UPSAMPLER_PATH = hf_hub_download(BASE_REPO, SPATIAL_UPSAMPLER_FILE, token=_HF_TOKEN) | |
| logger.info("Downloading distilled LoRA ...") | |
| DISTILLED_LORA_PATH = hf_hub_download(BASE_REPO, DISTILLED_LORA_FILE, token=_HF_TOKEN) | |
| logger.info("Downloading Relight IC-LoRA ...") | |
| LORA_PATH = hf_hub_download(LORA_REPO, LORA_FILE, token=_HF_TOKEN) | |
| logger.info("Downloading Gemma-3-12B text encoder (~24GB) ...") | |
| GEMMA_ROOT = snapshot_download( | |
| GEMMA_REPO, | |
| allow_patterns=["*.safetensors", "*.json", "tokenizer.model", "tokenizer.json"], | |
| token=_HF_TOKEN, | |
| ) | |
| logger.info("All weights ready.") | |
| # ----------------------------------------------------------------------------- | |
| # Build the IC-LoRA pipeline at module scope | |
| # ----------------------------------------------------------------------------- | |
| QUANTIZATION = build_fp8_cast_policy(CHECKPOINT_PATH) | |
| # The workflow uses TWO LoRAs: the Relight IC-LoRA at strength 1.0, and the | |
| # distilled LoRA at strength 0.5. Both are fused into the stage-1 transformer. | |
| LORAS = [ | |
| LoraPathStrengthAndSDOps(LORA_PATH, 1.0, LTXV_LORA_COMFY_RENAMING_MAP), | |
| LoraPathStrengthAndSDOps(DISTILLED_LORA_PATH, 0.5, LTXV_LORA_COMFY_RENAMING_MAP), | |
| ] | |
| PIPELINE = ICLoraPipeline( | |
| distilled_checkpoint_path=CHECKPOINT_PATH, | |
| spatial_upsampler_path=UPSAMPLER_PATH, | |
| gemma_root=GEMMA_ROOT, | |
| loras=LORAS, | |
| device=DEVICE, | |
| quantization=QUANTIZATION, | |
| ) | |
| def _preload_models_for_zerogpu(): | |
| """Build + cache the models at MODULE SCOPE so ZeroGPU preloads them.""" | |
| print("=" * 80, flush=True) | |
| print("Preloading models at module scope for ZeroGPU (.to('cuda') hijack)...", flush=True) | |
| print("=" * 80, flush=True) | |
| PIPELINE.prompt_encoder.warmup() | |
| PIPELINE.image_conditioner.warmup() | |
| PIPELINE.stage_1.warmup() | |
| PIPELINE.video_decoder.warmup() | |
| PIPELINE.audio_decoder.warmup() | |
| _maybe_preload_text_encoder() | |
| print("Models preloaded: weights registered for ZeroGPU streaming.", flush=True) | |
| PRELOAD_TEXT_ENCODER = os.environ.get("PRELOAD_TEXT_ENCODER", "1") == "1" | |
| def _maybe_preload_text_encoder(): | |
| if not PRELOAD_TEXT_ENCODER: | |
| print("[preload] text-encoder preload disabled; Gemma stays build-use-free.", flush=True) | |
| return | |
| from contextlib import nullcontext | |
| pe = PIPELINE.prompt_encoder | |
| try: | |
| if getattr(pe, "_cached_text_encoder", None) is None: | |
| pe._cached_text_encoder = pe._build_text_encoder() | |
| pe._text_encoder_ctx = lambda: nullcontext(pe._cached_text_encoder) | |
| print("[preload] Gemma text encoder preloaded resident for ZeroGPU.", flush=True) | |
| except Exception as e: | |
| pe._cached_text_encoder = None | |
| print(f"[preload] text-encoder preload skipped ({e!r}); using build-use-free.", flush=True) | |
| # ----------------------------------------------------------------------------- | |
| # Inference | |
| # ----------------------------------------------------------------------------- | |
| def _run(video_path, prompt, negative_prompt, seed, num_inference_steps, | |
| conditioning_strength, duration, azimuth, elevation, hardness, | |
| ball_state=None): | |
| num_frames = _plan_num_frames(video_path, duration) | |
| logger.info("Relight geometry: %d frames (%.2fs) %dx%d @ %dfps", | |
| num_frames, num_frames / TARGET_FPS, TARGET_WIDTH, TARGET_HEIGHT, TARGET_FPS) | |
| # Composite the light-direction sphere into every frame → reference video. | |
| # Prefer the snapshot from the interactive drag widget; fall back to the | |
| # pure-Python render for headless / API calls. | |
| ref_path = _composite_sphere_into_frames( | |
| video_path, num_frames, azimuth, elevation, hardness, ball_state | |
| ) | |
| sigmas = DISTILLED_SIGMAS | |
| if int(num_inference_steps) != (len(DISTILLED_SIGMAS) - 1): | |
| steps = int(num_inference_steps) | |
| idx = torch.linspace(0, len(DISTILLED_SIGMAS) - 1, steps + 1) | |
| sigmas = torch.tensor( | |
| [float(DISTILLED_SIGMAS[int(round(i.item()))]) for i in idx] | |
| ) | |
| sigmas[-1] = 0.0 | |
| video_iter, _audio = PIPELINE( | |
| prompt=prompt, | |
| seed=int(seed), | |
| height=TARGET_HEIGHT, | |
| width=TARGET_WIDTH, | |
| num_frames=num_frames, | |
| frame_rate=TARGET_FPS, | |
| images=[], # IC-LoRA: no first-frame image conditioning | |
| video_conditioning=[(ref_path, float(conditioning_strength))], | |
| conditioning_attention_strength=float(conditioning_strength), | |
| skip_stage_2=True, # single-stage distilled | |
| stage_1_sigmas=sigmas, | |
| ) | |
| tiling_config = TilingConfig.default() | |
| video_chunks_number = get_video_chunks_number(num_frames, tiling_config) | |
| out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name | |
| encode_video( | |
| video=video_iter, | |
| fps=TARGET_FPS, | |
| audio=None, | |
| output_path=out_path, | |
| video_chunks_number=video_chunks_number, | |
| ) | |
| return out_path | |
| def generate(video_path, light_look, azimuth, elevation, hardness, duration, | |
| seed, num_inference_steps, conditioning_strength, | |
| negative_prompt, ball_state="", progress=gr.Progress(track_tqdm=True)): | |
| """Relight a video clip using the LTX-2.3 Relight IC-LoRA. | |
| Args: | |
| video_path: Path to the source video to relight. | |
| light_look: One of the 12 trained lighting "looks" (e.g. "hard directional | |
| sunlight"). Combined with the trigger phrase to form the prompt. | |
| azimuth: Light azimuth in degrees (0 = front, 90 = right, 180 = behind). | |
| elevation: Light elevation in degrees (0 = horizon, 90 = top). | |
| hardness: Light hardness: 1 = soft, 2 = medium, 3 = hard. | |
| duration: Target output length in seconds (0.5–5.0). | |
| seed: RNG seed for reproducibility. | |
| num_inference_steps: Distilled denoising steps (8 is default). | |
| conditioning_strength: How strongly the reference guides the output. | |
| negative_prompt: Things to avoid in the output. | |
| ball_state: JSON string from the interactive drag-to-set light widget, | |
| carrying the rendered light-direction ball snapshot (PNG data-URL). | |
| When present it is the control signal composited into every frame; | |
| otherwise the demo falls back to the pure-Python sphere render. | |
| """ | |
| if not video_path: | |
| raise gr.Error("Please provide an input video to relight.") | |
| if not light_look: | |
| raise gr.Error("Please select a lighting look.") | |
| # If the interactive widget synced its state, prefer its azimuth/elevation/ | |
| # hardness so the composited snapshot and the (fallback) render agree. | |
| if ball_state: | |
| try: | |
| _st = json.loads(ball_state) | |
| if isinstance(_st, dict): | |
| azimuth = _st.get("azimuth", azimuth) | |
| elevation = _st.get("elevation", elevation) | |
| hardness = _st.get("hardness", hardness) | |
| except (TypeError, ValueError): | |
| pass | |
| # Build the prompt: trigger phrase + look + direction from the light ball | |
| prompt = f"{TRIGGER_PHRASE} {light_look}" | |
| return _run( | |
| video_path=video_path, | |
| prompt=prompt, | |
| negative_prompt=negative_prompt or DEFAULT_NEGATIVE, | |
| seed=seed, | |
| num_inference_steps=num_inference_steps, | |
| conditioning_strength=conditioning_strength, | |
| duration=duration, | |
| azimuth=azimuth, | |
| elevation=elevation, | |
| hardness=hardness, | |
| ball_state=ball_state, | |
| ) | |
| def update_preview(video_path, azimuth, elevation, hardness, ball_state=""): | |
| """Update the sphere composite preview when inputs change.""" | |
| if not video_path: | |
| return None | |
| try: | |
| return _preview_composite(video_path, azimuth, elevation, hardness, ball_state) | |
| except Exception as e: | |
| logger.warning("Preview failed: %s", e) | |
| return None | |
| # ----------------------------------------------------------------------------- | |
| # Interactive drag-to-set light-direction ball (Three.js widget) | |
| # ----------------------------------------------------------------------------- | |
| # Adapted from the custom HTML/JS component in | |
| # huggingface.co/spaces/linoyts/sun-direction-flux2-klein. The user clicks and | |
| # drags the ball to set the light direction; on release the widget renders the | |
| # lit matte sphere (light marker hidden) to a 512x512 PNG and writes | |
| # {"azimuth", "elevation", "hardness", "png": data-URL} into a hidden textbox | |
| # (#ball-state). That snapshot is the control signal composited into every frame. | |
| # The scene replicates the model card's sphere-render convention: | |
| # grey background RGB(173,173,173), matte sphere, camera (0,6,8) FOV 35°, | |
| # directional light from az/el, shadow tail pointing away from the light. | |
| LIGHT_BALL_HTML = """ | |
| <div id="light-ball" style="width:100%;max-width:260px;margin:0 auto;"> | |
| <div id="light-ball-mount" style="width:100%;aspect-ratio:1;border-radius:12px;overflow:hidden; | |
| cursor:grab;touch-action:none;background:#adadad;position:relative;"> | |
| <div id="light-ball-loading" style="position:absolute;inset:0;display:flex;align-items:center; | |
| justify-content:center;color:#555;font-family:sans-serif;font-size:0.85em;">loading light ball…</div> | |
| </div> | |
| <div id="light-ball-readout" style="text-align:center;font-family:monospace;font-size:0.8em; | |
| opacity:0.75;padding-top:6px;">💡 drag me · az 150° · el 15°</div> | |
| </div> | |
| """ | |
| LIGHT_BALL_JS = r""" | |
| <script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script> | |
| <script> | |
| (function () { | |
| if (window.__lightBallInited) return; | |
| window.__lightBallInited = true; | |
| // Defaults match the workflow's SphereLightNode (az 150, el 15, medium hardness). | |
| var az = 150, el = 15, hardness = 2; | |
| var HARDNESS_TO_INTENSITY = { 1: 0.8, 2: 1.5, 3: 2.5 }; | |
| function setGradioValue(elemId, value) { | |
| var container = document.getElementById(elemId); | |
| if (!container) return; | |
| var elInput = container.querySelector("input, textarea"); | |
| if (!elInput) return; | |
| var proto = elInput.tagName === "TEXTAREA" ? HTMLTextAreaElement.prototype : HTMLInputElement.prototype; | |
| Object.getOwnPropertyDescriptor(proto, "value").set.call(elInput, value); | |
| elInput.dispatchEvent(new Event("input", { bubbles: true })); | |
| elInput.dispatchEvent(new Event("change", { bubbles: true })); | |
| } | |
| function init() { | |
| var mount = document.getElementById("light-ball-mount"); | |
| if (!mount || !window.THREE) { setTimeout(init, 80); return; } | |
| var loading = document.getElementById("light-ball-loading"); | |
| if (loading) loading.remove(); | |
| var R = window.THREE; | |
| var canvas = document.createElement("canvas"); | |
| canvas.width = 512; canvas.height = 512; | |
| canvas.style.width = "100%"; canvas.style.height = "100%"; canvas.style.display = "block"; | |
| mount.appendChild(canvas); | |
| var renderer = new R.WebGLRenderer({ canvas: canvas, antialias: true, preserveDrawingBuffer: true }); | |
| renderer.shadowMap.enabled = true; | |
| renderer.shadowMap.type = R.PCFSoftShadowMap; | |
| renderer.setSize(512, 512, false); | |
| renderer.setClearColor(0xadadad); | |
| renderer.outputEncoding = R.sRGBEncoding; | |
| var scene = new R.Scene(); | |
| scene.background = new R.Color(0xadadad); | |
| var camera = new R.PerspectiveCamera(35, 1, 0.1, 200); | |
| camera.position.set(0, 6, 8); | |
| camera.lookAt(0, -0.5, 0); | |
| var plane = new R.Mesh(new R.PlaneGeometry(100, 100), | |
| new R.MeshStandardMaterial({ color: 0xadadad, roughness: 1, metalness: 0 })); | |
| plane.rotation.x = -Math.PI / 2; plane.position.y = -1; plane.receiveShadow = true; | |
| scene.add(plane); | |
| var sphere = new R.Mesh(new R.SphereGeometry(1, 64, 64), | |
| new R.MeshStandardMaterial({ color: 0xcccccc, roughness: 0.8, metalness: 0 })); | |
| sphere.castShadow = true; sphere.receiveShadow = true; | |
| scene.add(sphere); | |
| scene.add(new R.AmbientLight(0xffffff, 0.2)); | |
| var dirLight = new R.DirectionalLight(0xffffff, 1.5); | |
| dirLight.castShadow = true; | |
| dirLight.shadow.mapSize.width = 2048; dirLight.shadow.mapSize.height = 2048; | |
| dirLight.shadow.camera.near = 0.1; dirLight.shadow.camera.far = 50; | |
| dirLight.shadow.camera.left = -8; dirLight.shadow.camera.right = 8; | |
| dirLight.shadow.camera.top = 8; dirLight.shadow.camera.bottom = -8; | |
| dirLight.shadow.bias = -0.0005; dirLight.shadow.radius = 2; | |
| scene.add(dirLight); | |
| // Light marker: interactive view only — hidden before exporting the snapshot. | |
| var marker = new R.Mesh(new R.SphereGeometry(0.18, 24, 24), | |
| new R.MeshBasicMaterial({ color: 0xffcc33 })); | |
| scene.add(marker); | |
| function place() { | |
| var a = az * Math.PI / 180, e = el * Math.PI / 180, r = 10; | |
| dirLight.position.set(r * Math.cos(e) * Math.sin(a), r * Math.sin(e), r * Math.cos(e) * Math.cos(a)); | |
| dirLight.intensity = HARDNESS_TO_INTENSITY[hardness] || 1.5; | |
| marker.position.set(3.2 * Math.cos(e) * Math.sin(a), 3.2 * Math.sin(e), 3.2 * Math.cos(e) * Math.cos(a)); | |
| var ro = document.getElementById("light-ball-readout"); | |
| if (ro) ro.textContent = "💡 az " + Math.round(az) + "° · el " + Math.round(el) + "°"; | |
| } | |
| function render() { | |
| renderer.shadowMap.needsUpdate = true; | |
| renderer.render(scene, camera); | |
| } | |
| function syncState() { | |
| marker.visible = false; | |
| render(); | |
| var png = canvas.toDataURL("image/png"); | |
| marker.visible = true; | |
| render(); | |
| setGradioValue("ball-state", JSON.stringify({ azimuth: az, elevation: el, hardness: hardness, png: png })); | |
| // Keep the (hidden fallback) gradio sliders in sync with the widget. | |
| setGradioValue("ball-azimuth", Math.round(az)); | |
| setGradioValue("ball-elevation", Math.round(el)); | |
| } | |
| var dragging = false, lastX = 0, lastY = 0, moved = false; | |
| canvas.addEventListener("pointerdown", function (ev) { | |
| dragging = true; moved = false; lastX = ev.clientX; lastY = ev.clientY; | |
| canvas.setPointerCapture(ev.pointerId); | |
| canvas.style.cursor = "grabbing"; | |
| ev.preventDefault(); | |
| }); | |
| canvas.addEventListener("pointermove", function (ev) { | |
| if (!dragging) return; | |
| var dx = ev.clientX - lastX, dy = ev.clientY - lastY; | |
| lastX = ev.clientX; lastY = ev.clientY; | |
| if (Math.abs(dx) + Math.abs(dy) > 0) moved = true; | |
| az += dx * 0.6; | |
| az = ((az % 360) + 360) % 360; // wrap to [0, 360) | |
| el = Math.min(90, Math.max(-30, el - dy * 0.4)); // clamp to slider range | |
| place(); render(); | |
| }); | |
| function endDrag() { | |
| if (!dragging) return; | |
| dragging = false; | |
| canvas.style.cursor = "grab"; | |
| if (!moved) return; | |
| syncState(); | |
| } | |
| canvas.addEventListener("pointerup", endDrag); | |
| canvas.addEventListener("pointercancel", endDrag); | |
| // Hardness radio drives the light intensity live. | |
| function hookHardness() { | |
| var box = document.getElementById("ball-hardness"); | |
| if (!box) { setTimeout(hookHardness, 300); return; } | |
| box.querySelectorAll("input[type=radio]").forEach(function (inp) { | |
| inp.addEventListener("change", function () { | |
| if (inp.checked) { | |
| var v = parseInt(inp.value, 10); | |
| if (!isNaN(v)) { hardness = v; place(); render(); syncState(); } | |
| } | |
| }); | |
| }); | |
| } | |
| hookHardness(); | |
| place(); render(); | |
| setTimeout(syncState, 300); // initial snapshot so Generate works before any drag | |
| } | |
| init(); | |
| })(); | |
| </script> | |
| """ | |
| # ----------------------------------------------------------------------------- | |
| # UI | |
| # ----------------------------------------------------------------------------- | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| .hide-box { display: none !important; } | |
| """ | |
| with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, head=LIGHT_BALL_JS) as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """ | |
| # 💡 LTX-2.3 Relight — relight any video with a light-direction ball | |
| Relight an exterior video clip by choosing a light direction and a | |
| lighting style. Powered by | |
| [LTX-2.3-22B](https://huggingface.co/Lightricks/LTX-2.3) + the | |
| [Relight IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Relight). | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| video_in = gr.Video(label="Source video (exterior clip to relight)") | |
| with gr.Accordion("Light direction (drag the ball)", open=True): | |
| gr.Markdown( | |
| "**Click and drag the ball** to set the light direction. " | |
| "On release, the widget renders the light-direction ball " | |
| "in that position and composites it into every frame as " | |
| "the control signal." | |
| ) | |
| with gr.Row(): | |
| gr.HTML(LIGHT_BALL_HTML) | |
| composite_preview = gr.Image( | |
| label="Composite preview (first frame)", | |
| type="pil", | |
| height=260, | |
| interactive=False, | |
| ) | |
| hardness = gr.Radio( | |
| choices=[1, 2, 3], | |
| value=DEFAULT_HARDNESS, | |
| label="Light hardness", | |
| info="1 = soft, 2 = medium, 3 = hard directional light", | |
| elem_id="ball-hardness", | |
| ) | |
| # Fallback / manual azimuth+elevation. The drag widget keeps | |
| # these in sync (and headless API callers can set them | |
| # directly); they drive the pure-Python render when no | |
| # interactive snapshot is available. | |
| with gr.Accordion("Fine-tune direction (numeric)", open=False): | |
| azimuth = gr.Slider( | |
| 0, 360, value=DEFAULT_AZIMUTH, step=1, | |
| label="Azimuth (°)", | |
| info="0 = front, 90 = right, 180 = behind, 270 = left", | |
| elem_id="ball-azimuth", | |
| ) | |
| elevation = gr.Slider( | |
| -30, 90, value=DEFAULT_ELEVATION, step=1, | |
| label="Elevation (°)", | |
| info="0 = horizon, 90 = directly overhead", | |
| elem_id="ball-elevation", | |
| ) | |
| light_look = gr.Dropdown( | |
| choices=LIGHT_LOOKS, | |
| value=DEFAULT_LIGHT_LOOK, | |
| label="Lighting look", | |
| info="One of 12 trained lighting styles. Combined with the " | |
| "trigger phrase to form the prompt.", | |
| ) | |
| duration = gr.Slider( | |
| MIN_DURATION, MAX_DURATION, value=DEFAULT_DURATION, step=0.5, | |
| label="Video duration (seconds)", | |
| info="Output length (up to 5 s), trimmed to the source clip.", | |
| ) | |
| run = gr.Button("Relight video", variant="primary") | |
| with gr.Accordion("Advanced options", open=False): | |
| negative_prompt = gr.Textbox( | |
| label="Negative prompt", value=DEFAULT_NEGATIVE, lines=3 | |
| ) | |
| steps = gr.Slider( | |
| 4, 12, value=8, step=1, label="Denoising steps" | |
| ) | |
| conditioning_strength = gr.Slider( | |
| 0.0, 1.0, value=1.0, step=0.05, | |
| label="Reference conditioning strength", | |
| ) | |
| seed = gr.Number(value=42, precision=0, label="Seed") | |
| with gr.Column(): | |
| video_out = gr.Video(label="Relit video", autoplay=True) | |
| # Hidden wire: the interactive drag widget writes its JSON state here | |
| # ({"azimuth","elevation","hardness","png": data-URL}). The PNG is the | |
| # rendered light-direction ball snapshot used as the control signal. | |
| ball_state = gr.Textbox( | |
| value="", elem_id="ball-state", elem_classes=["hide-box"], visible=False | |
| ) | |
| # Update the first-frame composite preview when the direction/video change. | |
| ball_state.change( | |
| update_preview, | |
| inputs=[video_in, azimuth, elevation, hardness, ball_state], | |
| outputs=composite_preview, | |
| ) | |
| video_in.change( | |
| update_preview, | |
| inputs=[video_in, azimuth, elevation, hardness, ball_state], | |
| outputs=composite_preview, | |
| ) | |
| azimuth.input( | |
| update_preview, | |
| inputs=[video_in, azimuth, elevation, hardness, ball_state], | |
| outputs=composite_preview, | |
| ) | |
| elevation.input( | |
| update_preview, | |
| inputs=[video_in, azimuth, elevation, hardness, ball_state], | |
| outputs=composite_preview, | |
| ) | |
| hardness.change( | |
| update_preview, | |
| inputs=[video_in, azimuth, elevation, hardness, ball_state], | |
| outputs=composite_preview, | |
| ) | |
| inputs = [video_in, light_look, azimuth, elevation, hardness, duration, | |
| seed, steps, conditioning_strength, negative_prompt, ball_state] | |
| run.click(generate, inputs=inputs, outputs=video_out, api_name="generate") | |
| gr.Examples( | |
| examples=[ | |
| ["examples/street_scene.mp4", "hard directional sunlight", | |
| DEFAULT_AZIMUTH, DEFAULT_ELEVATION, DEFAULT_HARDNESS, | |
| DEFAULT_DURATION, 42, 8, 1.0], | |
| ["examples/people_running.mp4", "warm golden low front sun", | |
| 200.0, 25.0, 2, DEFAULT_DURATION, 42, 8, 1.0], | |
| ], | |
| inputs=[video_in, light_look, azimuth, elevation, hardness, | |
| duration, seed, steps, conditioning_strength], | |
| outputs=video_out, | |
| fn=lambda v, look, az, el, hard, dur, s, st, cs: generate( | |
| v, look, az, el, hard, dur, s, st, cs, DEFAULT_NEGATIVE, | |
| ), | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| # Build + pack the models now, in the MAIN process, so ZeroGPU preloads them. | |
| _preload_models_for_zerogpu() | |
| if __name__ == "__main__": | |
| demo.queue().launch(mcp_server=True, show_error=True) |