multimodalart's picture
multimodalart HF Staff
Fix: pass drag-widget JS via gr.Blocks(head=...) not launch()
47929ce verified
Raw
History Blame Contribute Delete
46.7 kB
"""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
# -----------------------------------------------------------------------------
@torch.inference_mode()
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
@spaces.GPU(duration=300, size="xlarge")
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)