linoyts HF Staff commited on
Commit
303ae8d
·
1 Parent(s): 6ecf85c

Switch backend to native LTX-2 (ICLoraPipeline) (#2)

Browse files

- Switch backend to native LTX-2 (ICLoraPipeline) (836c1bf35fbbc757a7c3eaf2ce909e3d10b4cda4)
- Hide UI controls the native distilled backend can't honor (751084892d27054a753b1bc32b4463f9a522cd4c)
- Revert README (keep name + card unchanged) (7bc1255e177901efe115419d194e54da1a5968d7)
- De-diffuser README (native backend; backlinks -> Lightricks/LTX-2.3) (58f659712370c2b0fd6ece9178eda06a4892f5d0)

Files changed (3) hide show
  1. README.md +2 -2
  2. app.py +214 -118
  3. requirements.txt +9 -7
README.md CHANGED
@@ -11,10 +11,10 @@ pinned: false
11
  hardware: zero-a10g
12
  short_description: Add water VFX to video with an LTX-2.3 IC-LoRA
13
  models:
14
- - diffusers/LTX-2.3-Distilled-Diffusers
15
  - Lightricks/LTX-2.3-22b-IC-LoRA-Water-Simulation
16
  ---
17
 
18
  # 🌊 LTX-2.3 Water Simulation
19
  Adds naturally-moving water (rivers, surf, rain, floods, splashes) to a dry clip while preserving subject,
20
- framing and camera. IC-LoRA on distilled LTX-2.3 (`LTX2InContextPipeline`, 8-step, `ADD WATER` trigger, strength ~1.2).
 
11
  hardware: zero-a10g
12
  short_description: Add water VFX to video with an LTX-2.3 IC-LoRA
13
  models:
14
+ - Lightricks/LTX-2.3
15
  - Lightricks/LTX-2.3-22b-IC-LoRA-Water-Simulation
16
  ---
17
 
18
  # 🌊 LTX-2.3 Water Simulation
19
  Adds naturally-moving water (rivers, surf, rain, floods, splashes) to a dry clip while preserving subject,
20
+ framing and camera. IC-LoRA on distilled LTX-2.3 (the native LTX-2 pipeline, 8-step, `ADD WATER` trigger, strength ~1.2).
app.py CHANGED
@@ -1,65 +1,137 @@
1
  import os
 
 
2
 
3
- os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
4
- os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
 
5
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6
  import random
7
  import tempfile
8
- import threading
9
- import time
10
 
11
  import numpy as np
12
  import imageio.v3 as iio
13
- import spaces
 
14
  import torch
 
 
 
 
15
  import gradio as gr
16
- from PIL import Image, ImageOps
17
- from huggingface_hub import hf_hub_download
18
- from safetensors.torch import load_file
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19
 
20
- from diffusers import LTX2InContextPipeline
21
- from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
22
- from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
23
- from diffusers.utils import load_video, encode_video
24
 
25
- # --- Config -----------------------------------------------------------------
26
- # Water-simulation IC-LoRA — distilled recipe (8 sigmas, CFG off), strength sweet-spot ~1.2.
27
- BASE_MODEL = "diffusers/LTX-2.3-Distilled-Diffusers"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
28
  LORA_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-Water-Simulation"
29
  LORA_FILE = "ltx-2.3-22b-ic-lora-water-simulation-0.9.safetensors"
30
- LORA_SCALE = 1.2
31
- FPS = 24
32
- NUM_STEPS = len(DISTILLED_SIGMA_VALUES)
33
- MAX_SEED = np.iinfo(np.int32).max
34
- HF_TOKEN = os.environ.get("HF_TOKEN")
35
-
36
- # Card validates at 1920×1088; offer higher buckets (water under-renders at low res). Capped below 1920 for ZeroGPU speed.
37
- RES_PRESETS = {"960×544 (fast)": (960, 544), "1216×704 (recommended)": (1216, 704),
38
- "1536×864 (high)": (1536, 864), "1920×1088 (native)": (1920, 1088)}
39
  FRAME_CHOICES = [49, 73, 97, 121]
 
40
 
41
- pipe = LTX2InContextPipeline.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16)
42
- pipe.to("cuda")
43
- pipe.vae.enable_tiling()
44
- _lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)
45
- pipe.load_lora_weights(load_file(_lora_path), adapter_name="water")
46
- pipe.fuse_lora(lora_scale=LORA_SCALE)
47
- pipe.unload_lora_weights()
48
- # AOTI: load precompiled blocks at ROOT level. LoRA is fused+unloaded so weight names
49
- # match the compiled constants. Scale is per-request, so adjust by re-fusing the *delta*
50
- # (fuse_lora is additive) + re-snapshot, only when it changes.
51
- spaces.aoti_load(module=pipe.transformer, repo_id="ltx-community/LTX-2.3-Transformer-GroupA-sm120-cu130-r9e")
52
- _FUSED_SCALE = LORA_SCALE
53
- def _refuse(scale):
54
- global _FUSED_SCALE
55
- delta = scale - _FUSED_SCALE
56
- if delta == 0:
57
- return
58
- pipe.load_lora_weights(load_file(_lora_path), adapter_name="water")
59
- pipe.fuse_lora(lora_scale=delta)
60
- pipe.unload_lora_weights()
61
- spaces.aoti_load(module=pipe.transformer, repo_id="ltx-community/LTX-2.3-Transformer-GroupA-sm120-cu130-r9e")
62
- _FUSED_SCALE = scale
63
 
64
 
65
  def _src_fps(path, default=FPS):
@@ -69,94 +141,124 @@ def _src_fps(path, default=FPS):
69
  return default
70
 
71
 
72
- def _load_frames(path, num_frames, width, height):
73
- frames = load_video(path)
74
- if not frames:
75
- return []
76
- fps = _src_fps(path)
77
  out = []
78
  for i in range(num_frames):
79
- idx = min(int(round(i / FPS * fps)), len(frames) - 1)
80
- out.append(ImageOps.fit(frames[idx].convert("RGB"), (width, height), Image.LANCZOS))
81
- return out
 
 
 
 
 
 
82
 
83
 
84
- def _pick_resolution(first_frame, preset):
85
  w, h = RES_PRESETS[preset]
86
- if first_frame.height > first_frame.width:
87
- w, h = h, w
 
 
 
 
88
  return w, h
89
 
90
 
91
- def _build_prompt(prompt):
92
- desc = prompt.strip() or "a flowing stream of clear water"
93
- return (
94
- "Reference shows the dry scene. Edited shows the same scene with water added. "
95
- f"ADD WATER {desc}. "
96
- "Subject identity, clothing, framing, and background geometry are identical to the reference; "
97
- "only water-related elements differ between reference and edited."
98
- )
99
 
 
 
 
 
 
 
 
 
 
 
 
100
 
101
- def _export(video_np, audio, path):
102
- kw = {}
103
- if audio is not None:
104
- kw = dict(audio=audio[0].float().cpu(), audio_sample_rate=pipe.vocoder.config.output_sampling_rate)
105
- encode_video(video_np, fps=FPS, output_path=path, **kw)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
106
 
107
 
108
  def _duration(*args, **kwargs):
109
- preset = next((a for a in args if isinstance(a, str) and a in RES_PRESETS), "960×544 (fast)")
110
- num_frames = next((a for a in args if isinstance(a, int) and a in FRAME_CHOICES), 73)
111
- w, h = RES_PRESETS[preset]
112
- per_frame = max(1.0, (w * h) / (768 * 448)) # scale with pixel count
113
- return int(25 + int(num_frames) * per_frame * 0.65) # ~2.3x measured runtime
114
 
115
 
116
  @spaces.GPU(duration=_duration)
117
- def add_water(video, prompt, strength, preset, num_frames, seed, randomize,
118
- progress=gr.Progress(track_tqdm=True)):
119
  if video is None:
120
- raise gr.Error("Please upload a 'dry' video to add water to.")
121
- if randomize:
122
- seed = random.randint(0, MAX_SEED)
123
- seed = int(seed)
124
  num_frames = int(num_frames)
125
-
126
- probe = load_video(video)
127
- if not probe:
128
- raise gr.Error("Could not read any frames from that video.")
129
- width, height = _pick_resolution(probe[0], preset)
130
- ref = _load_frames(video, num_frames, width, height)
131
- _refuse(float(strength))
132
- full_prompt = _build_prompt(prompt)
133
-
134
- def _cb(p, i, t, kw):
135
- progress((i + 1) / NUM_STEPS, desc=f"Adding water — step {i + 1}/{NUM_STEPS}")
136
- return {}
137
-
138
- video_out, audio_out = pipe(
139
- prompt=full_prompt, negative_prompt="",
140
- reference_conditions=[LTX2ReferenceCondition(frames=ref, strength=1.0)],
141
- reference_downscale_factor=1,
142
- width=width, height=height, num_frames=num_frames, frame_rate=FPS,
143
- num_inference_steps=NUM_STEPS, sigmas=DISTILLED_SIGMA_VALUES,
144
- guidance_scale=1.0, stg_scale=0.0, audio_guidance_scale=1.0, audio_stg_scale=0.0,
145
- generator=torch.Generator(device="cuda").manual_seed(seed),
146
- output_type="np", return_dict=False,
147
- )
148
- out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
149
- _export(video_out[0], audio_out, out_path)
150
  return out_path, seed
151
 
152
 
 
 
 
 
 
 
 
153
  with gr.Blocks(title="LTX-2.3 Water Simulation") as demo:
154
  gr.Markdown(
155
  "# 🌊 LTX-2.3 Water Simulation\n"
156
  "Add believable, naturally-moving water to a dry clip — rivers, surf, rain, waterfalls, floods, "
157
  "splashes — that interacts with the moving scene, while maintaining subject and framing identity. "
158
- "Using [LTX 2.3 Distilled](https://huggingface.co/diffusers/LTX-2.3-Distilled-Diffusers) with the "
159
- "[Water Simulation IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Water-Simulation), via diffusers 🧨."
160
  )
161
  with gr.Row():
162
  with gr.Column():
@@ -166,8 +268,6 @@ with gr.Blocks(title="LTX-2.3 Water Simulation") as demo:
166
  placeholder="a clear shallow stream braiding around their legs with white foam crests and glistening wet ground; rushing water, gentle splashing",
167
  )
168
  with gr.Accordion("Settings", open=False):
169
- strength = gr.Slider(1.0, 1.6, value=1.2, step=0.05,
170
- label="Water strength (1.2–1.3 natural · 1.35+ hard surface→sea · ≥1.5 max drama)")
171
  preset = gr.Dropdown(list(RES_PRESETS), value="1216×704 (recommended)", label="Resolution")
172
  num_frames = gr.Dropdown(FRAME_CHOICES, value=73, label="Frames (24fps)")
173
  randomize = gr.Checkbox(True, label="Randomize seed")
@@ -176,19 +276,15 @@ with gr.Blocks(title="LTX-2.3 Water Simulation") as demo:
176
  with gr.Column():
177
  video_out = gr.Video(label="Result with water")
178
 
179
- run.click(add_water, inputs=[video_in, prompt, strength, preset, num_frames, seed, randomize],
180
  outputs=[video_out, seed])
181
 
182
  gr.Examples(
183
  examples=[
184
- ["examples/man_dancing_dry.mp4",
185
- "a clear, shallow stream rushing and braiding around their legscold mountain water swirling with white foam crests and glassy ripples, the wet floor glistening and throwing back reflections, bright droplets kicked up with every step; lively rushing water and rhythmic splashes as they move",
186
- 1.3, "1216×704 (recommended)", 73, 42, False],
187
- ["examples/landscape_dry.mp4",
188
- "a wide river flooding across the valley floor — the water spreading in glassy sheets with rippling reflections of the sky and drifting ribbons of white foam, a soft mist rising off the surface in the cool light; steadily flowing water and the distant rush of a waterfall",
189
- 1.25, "1216×704 (recommended)", 73, 42, False],
190
  ],
191
- inputs=[video_in, prompt, strength, preset, num_frames, seed, randomize],
192
  outputs=[video_out, seed], fn=add_water, cache_examples=True, cache_mode="lazy",
193
  )
194
 
 
1
  import os
2
+ import subprocess
3
+ import sys
4
 
5
+ # ZeroGPU: torch.compile / dynamo unsupported — disable before any torch import.
6
+ os.environ["TORCH_COMPILE_DISABLE"] = "1"
7
+ os.environ["TORCHDYNAMO_DISABLE"] = "1"
8
 
9
+ # memory-efficient attention
10
+ subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)
11
+
12
+ # --- clone + install the NATIVE LTX-2 codebase at the pinned commit the working ZeroGPU spaces use ---
13
+ LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
14
+ LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
15
+ LTX_COMMIT = "ae855f8538843825f9015a419cf4ba5edaf5eec2"
16
+ if not os.path.exists(LTX_REPO_DIR):
17
+ subprocess.run(["git", "clone", LTX_REPO_URL, LTX_REPO_DIR], check=True)
18
+ subprocess.run(["git", "-C", LTX_REPO_DIR, "checkout", LTX_COMMIT], check=True)
19
+ subprocess.run([sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
20
+ "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-core"),
21
+ "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")], check=True)
22
+ sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src"))
23
+ sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))
24
+
25
+ import logging
26
  import random
27
  import tempfile
 
 
28
 
29
  import numpy as np
30
  import imageio.v3 as iio
31
+ from PIL import Image, ImageOps
32
+
33
  import torch
34
+ torch._dynamo.config.suppress_errors = True
35
+ torch._dynamo.config.disable = True
36
+
37
+ import spaces
38
  import gradio as gr
39
+ from huggingface_hub import hf_hub_download, snapshot_download
40
+
41
+ # Import LTX modules in the proven order — importing ltx_core.quantization/loader FIRST hits a
42
+ # circular import (fp8_cast <-> loader.fuse_loras). Importing the model modules first forces the
43
+ # correct init order (mirrors the working reference Space).
44
+ from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number, decode_video as _vae_decode_video # noqa: F401
45
+ from ltx_core.model.upsampler import upsample_video as _upsample_video # noqa: F401
46
+ from ltx_core.model.audio_vae import encode_audio as _vae_encode_audio # noqa: F401
47
+ from ltx_core.quantization import QuantizationPolicy
48
+ from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
49
+ from ltx_pipelines.ic_lora import ICLoraPipeline
50
+ from ltx_pipelines.utils.media_io import encode_video
51
+
52
+ # --- ZeroGPU loader patch -------------------------------------------------------------
53
+ # The native loader opens safetensors directly on the CUDA device
54
+ # (safe_open(path, device="cuda")), doing the host->device copy in safetensors' own C++
55
+ # (cudaMemcpy) — bypassing torch.Tensor.to, the call ZeroGPU patches to virtualise + pack
56
+ # weights at module scope. Result: "No CUDA GPUs are available" at startup, nothing packs.
57
+ # Patch it to open on CPU then move via torch.Tensor.to (ZeroGPU-virtualisable).
58
+ import safetensors as _safetensors
59
+ import ltx_core.loader.sft_loader as _sft
60
+ from ltx_core.loader.primitives import StateDict as _StateDict
61
+
62
+ def _zerogpu_safe_load(self, path, sd_ops, device=None):
63
+ device = device or torch.device("cpu")
64
+ sd, size, dtype = {}, 0, set()
65
+ model_paths = path if isinstance(path, list) else [path]
66
+ for shard_path in model_paths:
67
+ with _safetensors.safe_open(shard_path, framework="pt", device="cpu") as f:
68
+ for name in f.keys():
69
+ expected = name if sd_ops is None else sd_ops.apply_to_key(name)
70
+ if expected is None:
71
+ continue
72
+ value = f.get_tensor(name).to(device=device) # torch path -> ZeroGPU-virtualised
73
+ kvs = ((expected, value),)
74
+ if sd_ops is not None:
75
+ kvs = sd_ops.apply_to_key_value(expected, value)
76
+ for k, v in kvs:
77
+ size += v.nbytes
78
+ dtype.add(v.dtype)
79
+ sd[k] = v
80
+ return _StateDict(sd=sd, device=device, size=size, dtype=dtype)
81
 
82
+ _sft.SafetensorsStateDictLoader.load = _zerogpu_safe_load
83
+ print("[PATCH] safetensors loader -> CPU-open + torch.to (ZeroGPU-virtualisable)")
84
+ # --------------------------------------------------------------------------------------
 
85
 
86
+ # --- attention backend patch (FA3 crashes on Blackwell ZeroGPU; use xformers/SDPA) ---
87
+ import torch.nn.functional as F
88
+ from ltx_core.model.transformer import attention as _attn_mod
89
+
90
+ def _sdpa_as_mea(query, key, value, attn_bias=None, scale=None, **kwargs):
91
+ q, k, v = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
92
+ return F.scaled_dot_product_attention(q, k, v, scale=scale).transpose(1, 2)
93
+
94
+ # IMPORTANT (ZeroGPU): never query CUDA at module scope. SDPA works on every GPU (incl.
95
+ # Blackwell ZeroGPU, where FA3 crashes), so patch it unconditionally.
96
+ _attn_mod.memory_efficient_attention = _sdpa_as_mea
97
+ print("[ATTN] SDPA (patched at module scope, no CUDA query)")
98
+
99
+ logging.getLogger().setLevel(logging.INFO)
100
+
101
+ # =========================== PER-LORA CONFIG (colorize) ===========================
102
+ TITLE = "LTX-2.3 Add Water (native LTX-2)"
103
  LORA_REPO = "Lightricks/LTX-2.3-22b-IC-LoRA-Water-Simulation"
104
  LORA_FILE = "ltx-2.3-22b-ic-lora-water-simulation-0.9.safetensors"
105
+ LORA_SCALE = 1.0
106
+ SKIP_STAGE_2 = True
107
+ GRAYSCALE_REF = False
108
+ RES_PRESETS = {"960×544 (fast)": (960, 544), "1216×704 (recommended)": (1216, 704)}
109
+ DEFAULT_PRESET = "1216×704 (recommended)"
 
 
 
 
110
  FRAME_CHOICES = [49, 73, 97, 121]
111
+ DEFAULT_FRAMES = 73
112
 
113
+ def build_prompt(p):
114
+ return (
115
+ "Reference shows the dry scene. Edited shows the same scene with realistic, naturally-moving water added. "
116
+ f"ADD WATER {p.strip()}. "
117
+ "Subject identity, framing and motion are identical to the reference; only water is added."
118
+ )
119
+
120
+ EXAMPLES = [
121
+ ["examples/landscape_dry.mp4",
122
+ "a wide river flooding across the valley with glassy rippling reflections and drifting foam, mist rising; flowing water and a distant waterfall",
123
+ "1216×704 (recommended)", 73, 42, False],
124
+ ["examples/man_dancing_dry.mp4",
125
+ "a clear shallow stream rushing and braiding around their legs with white foam and splashes; rushing water and rhythmic splashes",
126
+ "1216×704 (recommended)", 73, 42, False],
127
+ ]
128
+ # =================================================================================
129
+
130
+ FPS = 24.0
131
+ MAX_SEED = np.iinfo(np.int32).max
132
+ HF_TOKEN = os.environ.get("HF_TOKEN")
133
+ LTX_MODEL_REPO = "Lightricks/LTX-2.3"
134
+ GEMMA_REPO = "google/gemma-3-12b-it-qat-q4_0-unquantized"
135
 
136
 
137
  def _src_fps(path, default=FPS):
 
141
  return default
142
 
143
 
144
+ def _prep_reference(path, width, height, num_frames):
145
+ """Resample to 24fps, aspect-fit/crop to WxH, NF frames; (optionally grayscale); write temp mp4."""
146
+ vid = iio.imread(path, plugin="pyav")
147
+ src_fps = _src_fps(path)
148
+ n = len(vid)
149
  out = []
150
  for i in range(num_frames):
151
+ idx = min(int(round(i / FPS * src_fps)), n - 1)
152
+ im = Image.fromarray(vid[idx]).convert("RGB")
153
+ im = ImageOps.fit(im, (width, height), Image.LANCZOS)
154
+ if GRAYSCALE_REF:
155
+ im = im.convert("L").convert("RGB")
156
+ out.append(np.array(im))
157
+ tmp = tempfile.mktemp(suffix=".mp4")
158
+ iio.imwrite(tmp, np.stack(out), fps=FPS, plugin="pyav", codec="libx264")
159
+ return tmp
160
 
161
 
162
+ def _pick_resolution(path, preset):
163
  w, h = RES_PRESETS[preset]
164
+ try:
165
+ f0 = iio.imread(path, plugin="pyav", index=0)
166
+ if f0.shape[0] > f0.shape[1]: # portrait
167
+ w, h = h, w
168
+ except Exception:
169
+ pass
170
  return w, h
171
 
172
 
173
+ # --- Load native pipeline + IC-LoRA once at module scope (ZeroGPU packs weights here) ---
174
+ print("Downloading checkpoints…")
175
+ checkpoint_path = hf_hub_download(LTX_MODEL_REPO, "ltx-2.3-22b-distilled-1.1.safetensors", token=HF_TOKEN)
176
+ spatial_upsampler_path = hf_hub_download(LTX_MODEL_REPO, "ltx-2.3-spatial-upscaler-x2-1.1.safetensors", token=HF_TOKEN)
177
+ gemma_root = snapshot_download(GEMMA_REPO, token=HF_TOKEN)
178
+ lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=HF_TOKEN)
 
 
179
 
180
+ print("Building ICLoraPipeline…")
181
+ pipeline = ICLoraPipeline(
182
+ distilled_checkpoint_path=checkpoint_path,
183
+ spatial_upsampler_path=spatial_upsampler_path,
184
+ gemma_root=gemma_root,
185
+ loras=[LoraPathStrengthAndSDOps(lora_path, LORA_SCALE, LTXV_LORA_COMFY_RENAMING_MAP)],
186
+ # bf16 (NOT fp8): the IC-LoRA is fused into the transformer at MODULE SCOPE (the GPU
187
+ # worker can't re-open the checkpoint file). fp8_cast()'s fusion runs a custom CUDA kernel
188
+ # that can't be ZeroGPU-virtualised; the bf16 fuse rule is pure torch -> virtualisable.
189
+ quantization=None,
190
+ )
191
 
192
+
193
+ def _preload_pin(ledger, tag):
194
+ if ledger is None:
195
+ return
196
+ for name in ["transformer", "video_encoder", "video_decoder", "audio_encoder",
197
+ "audio_decoder", "vocoder", "spatial_upsampler", "text_encoder",
198
+ "gemma_embeddings_processor"]:
199
+ fn = getattr(ledger, name, None)
200
+ if callable(fn):
201
+ try:
202
+ obj = fn()
203
+ setattr(ledger, name, (lambda o=obj: o))
204
+ print(f"[preload {tag}] {name} ✓")
205
+ except Exception as e:
206
+ print(f"[preload {tag}] {name} skipped: {e}")
207
+
208
+ # Preload stage 1 always; preload stage 2 only when two-stage is used (skip_stage_2=False).
209
+ # Eagerly pinning both ledgers materializes TWO ~46GB transformers — too big for the ZeroGPU pack.
210
+ _preload_pin(getattr(pipeline, "stage_1_model_ledger", None), "stage1")
211
+ if not SKIP_STAGE_2:
212
+ _preload_pin(getattr(pipeline, "stage_2_model_ledger", None), "stage2")
213
+ print("Pipeline ready.")
214
 
215
 
216
  def _duration(*args, **kwargs):
217
+ nf = next((a for a in args if isinstance(a, int) and a in FRAME_CHOICES), DEFAULT_FRAMES)
218
+ return int(60 + nf * 1.2)
 
 
 
219
 
220
 
221
  @spaces.GPU(duration=_duration)
222
+ @torch.inference_mode()
223
+ def add_water(video, prompt, preset, num_frames, seed, randomize, progress=gr.Progress(track_tqdm=True)):
224
  if video is None:
225
+ raise gr.Error("Please upload a video.")
226
+ if not prompt.strip():
227
+ raise gr.Error("Describe the result (e.g. 'a brown rabbit on grey rocks, soft birdsong').")
228
+ seed = random.randint(0, MAX_SEED) if randomize else int(seed)
229
  num_frames = int(num_frames)
230
+ width, height = _pick_resolution(video, preset)
231
+ ref_path = _prep_reference(video, width, height, num_frames)
232
+ tiling = TilingConfig.default()
233
+ # skip_stage_2 outputs at half the passed dims -> pass 2x so output matches the preset.
234
+ gen_w, gen_h = (width * 2, height * 2) if SKIP_STAGE_2 else (width, height)
235
+ video_out, audio_out = pipeline(
236
+ prompt=build_prompt(prompt),
237
+ seed=seed, height=gen_h, width=gen_w,
238
+ num_frames=num_frames, frame_rate=FPS,
239
+ images=[], video_conditioning=[(ref_path, 1.0)],
240
+ skip_stage_2=SKIP_STAGE_2, tiling_config=tiling,
241
+ )
242
+ out_path = tempfile.mktemp(suffix=".mp4")
243
+ encode_video(video=video_out, fps=FPS, audio=audio_out, output_path=out_path,
244
+ video_chunks_number=get_video_chunks_number(num_frames, tiling))
 
 
 
 
 
 
 
 
 
 
245
  return out_path, seed
246
 
247
 
248
+ # --- UI config (match the public Space exactly) ---
249
+ RES_PRESETS = {"960×544 (fast)": (960, 544), "1216×704 (recommended)": (1216, 704),
250
+ "1536×864 (high)": (1536, 864), "1920×1088 (native)": (1920, 1088)}
251
+ FRAME_CHOICES = [49, 73, 97, 121]
252
+ LORA_SCALE = 1.2 # native fixed scale (matches public default)
253
+
254
+
255
  with gr.Blocks(title="LTX-2.3 Water Simulation") as demo:
256
  gr.Markdown(
257
  "# 🌊 LTX-2.3 Water Simulation\n"
258
  "Add believable, naturally-moving water to a dry clip — rivers, surf, rain, waterfalls, floods, "
259
  "splashes — that interacts with the moving scene, while maintaining subject and framing identity. "
260
+ "Using [LTX 2.3 Distilled](https://huggingface.co/Lightricks/LTX-2.3) with the "
261
+ "[Water Simulation IC-LoRA](https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-Water-Simulation)."
262
  )
263
  with gr.Row():
264
  with gr.Column():
 
268
  placeholder="a clear shallow stream braiding around their legs with white foam crests and glistening wet ground; rushing water, gentle splashing",
269
  )
270
  with gr.Accordion("Settings", open=False):
 
 
271
  preset = gr.Dropdown(list(RES_PRESETS), value="1216×704 (recommended)", label="Resolution")
272
  num_frames = gr.Dropdown(FRAME_CHOICES, value=73, label="Frames (24fps)")
273
  randomize = gr.Checkbox(True, label="Randomize seed")
 
276
  with gr.Column():
277
  video_out = gr.Video(label="Result with water")
278
 
279
+ run.click(add_water, inputs=[video_in, prompt, preset, num_frames, seed, randomize],
280
  outputs=[video_out, seed])
281
 
282
  gr.Examples(
283
  examples=[
284
+ ['examples/man_dancing_dry.mp4', 'a clear, shallow stream rushing and braiding around their legs — cold mountain water swirling with white foam crests and glassy ripples, the wet floor glistening and throwing back reflections, bright droplets kicked up with every step; lively rushing water and rhythmic splashes as they move', '1216×704 (recommended)', 73, 42, False],
285
+ ['examples/landscape_dry.mp4', 'a wide river flooding across the valley floorthe water spreading in glassy sheets with rippling reflections of the sky and drifting ribbons of white foam, a soft mist rising off the surface in the cool light; steadily flowing water and the distant rush of a waterfall', '1216×704 (recommended)', 73, 42, False],
 
 
 
 
286
  ],
287
+ inputs=[video_in, prompt, preset, num_frames, seed, randomize],
288
  outputs=[video_out, seed], fn=add_water, cache_examples=True, cache_mode="lazy",
289
  )
290
 
requirements.txt CHANGED
@@ -1,9 +1,11 @@
1
- git+https://github.com/huggingface/diffusers
2
- transformers
3
  accelerate
4
- peft
5
- safetensors
6
- sentencepiece
7
- imageio
8
- imageio-ffmpeg
9
  av
 
 
 
 
 
1
+ transformers==4.57.6
 
2
  accelerate
3
+ torch==2.8.0
4
+ torchaudio==2.8.0
5
+ einops
6
+ scipy
 
7
  av
8
+ scikit-image>=0.25.2
9
+ flashpack==0.1.2
10
+ imageio[ffmpeg]
11
+ pillow