dagloop5 commited on
Commit
4b5cf9a
·
verified ·
1 Parent(s): 4a7b271

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -938
app.py DELETED
@@ -1,938 +0,0 @@
1
- # =============================================================================
2
- # Installation and Setup
3
- # =============================================================================
4
- import os
5
- import subprocess
6
- import sys
7
-
8
- os.environ["TORCH_COMPILE_DISABLE"] = "1"
9
- os.environ["TORCHDYNAMO_DISABLE"] = "1"
10
-
11
- subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)
12
-
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
- # This commit has TI2VidTwoStagesHQPipeline with ModelLedger
16
- LTX_COMMIT = "ae855f8538843825f9015a419cf4ba5edaf5eec2"
17
-
18
- if not os.path.exists(LTX_REPO_DIR):
19
- print(f"Cloning {LTX_REPO_URL}...")
20
- subprocess.run(["git", "clone", LTX_REPO_URL, LTX_REPO_DIR], check=True)
21
- subprocess.run(["git", "checkout", LTX_COMMIT], cwd=LTX_REPO_DIR, check=True)
22
-
23
- print("Installing ltx-core and ltx-pipelines from cloned repo...")
24
- subprocess.run(
25
- [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "-e",
26
- os.path.join(LTX_REPO_DIR, "packages", "ltx-core"),
27
- "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")],
28
- check=True,
29
- )
30
-
31
- sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src"))
32
- sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))
33
-
34
- import logging
35
- import random
36
- import tempfile
37
- from pathlib import Path
38
- import gc
39
- import hashlib
40
-
41
- import torch
42
- torch._dynamo.config.suppress_errors = True
43
- torch._dynamo.config.disable = True
44
-
45
- import spaces
46
- import gradio as gr
47
- import numpy as np
48
- from huggingface_hub import hf_hub_download, snapshot_download
49
- from safetensors.torch import load_file, save_file
50
-
51
- from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
52
- from ltx_core.quantization import QuantizationPolicy
53
- from ltx_core.loader import LoraPathStrengthAndSDOps, LTXV_LORA_COMFY_RENAMING_MAP
54
- from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
55
- from ltx_core.components.noisers import GaussianNoiser
56
- from ltx_core.components.diffusion_steps import Res2sDiffusionStep
57
- from ltx_core.components.schedulers import LTX2Scheduler
58
- from ltx_core.types import Audio, LatentState, VideoPixelShape
59
- from ltx_core.tools import VideoLatentShape
60
- from ltx_pipelines.ti2vid_two_stages_hq import TI2VidTwoStagesHQPipeline
61
- from ltx_pipelines.utils.args import ImageConditioningInput
62
- from ltx_pipelines.utils.constants import LTX_2_3_HQ_PARAMS, STAGE_2_DISTILLED_SIGMA_VALUES
63
- from ltx_pipelines.utils.media_io import encode_video
64
- from ltx_pipelines.utils.helpers import (
65
- assert_resolution,
66
- cleanup_memory,
67
- combined_image_conditionings,
68
- get_device,
69
- )
70
-
71
- # Patch xformers
72
- try:
73
- from ltx_core.model.transformer import attention as _attn_mod
74
- from xformers.ops import memory_efficient_attention as _mea
75
- _attn_mod.memory_efficient_attention = _mea
76
- print("[ATTN] xformers patch applied")
77
- except Exception as e:
78
- print(f"[ATTN] xformers patch failed: {e}")
79
-
80
- logging.getLogger().setLevel(logging.INFO)
81
-
82
- MAX_SEED = np.iinfo(np.int32).max
83
- DEFAULT_PROMPT = (
84
- "A majestic eagle soaring over mountain peaks at sunset, "
85
- "wings spread wide against the orange sky, feathers catching the light, "
86
- "wind currents visible in the motion blur, cinematic slow motion, 4K quality"
87
- )
88
- DEFAULT_NEGATIVE_PROMPT = (
89
- "worst quality, inconsistent motion, blurry, jittery, distorted, "
90
- "deformed, artifacts, text, watermark, logo, frame, border, "
91
- "low resolution, pixelated, unnatural, fake, CGI, cartoon"
92
- )
93
- DEFAULT_FRAME_RATE = 24.0
94
- MIN_DIM, MAX_DIM, STEP = 256, 1280, 64
95
- MIN_FRAMES, MAX_FRAMES = 9, 257
96
-
97
- RESOLUTIONS = {
98
- "16:9": {"width": 1280, "height": 704},
99
- "9:16": {"width": 704, "height": 1280},
100
- "1:1": {"width": 960, "height": 960},
101
- }
102
-
103
- LTX_MODEL_REPO = "Lightricks/LTX-2.3"
104
- GEMMA_REPO = "Lightricks/gemma-3-12b-it-qat-q4_0-unquantized"
105
-
106
- # =============================================================================
107
- # Custom HQ Pipeline with LoRA Cache Support
108
- # =============================================================================
109
-
110
- class HQPipelineWithCachedLoRA(TI2VidTwoStagesHQPipeline):
111
- """
112
- TI2VidTwoStagesHQPipeline with support for cached LoRA state.
113
- Bypasses GPU LoRA fusion by applying pre-cached state_dict.
114
- Supports negative prompts and guidance parameters.
115
- """
116
-
117
- def __init__(self, *args, **kwargs):
118
- super().__init__(*args, **kwargs)
119
- # Storage for cached LoRA states for each stage
120
- self._cached_state_stage1 = None
121
- self._cached_state_stage2 = None
122
-
123
- def apply_cached_lora_state(self, state_dict_stage1, state_dict_stage2=None):
124
- """Apply pre-cached LoRA state to both stage transformers."""
125
- self._cached_state_stage1 = state_dict_stage1
126
- self._cached_state_stage2 = state_dict_stage2 if state_dict_stage2 else state_dict_stage1
127
-
128
- @torch.inference_mode()
129
- def __call__( # noqa: PLR0913
130
- self,
131
- prompt: str,
132
- negative_prompt: str,
133
- seed: int,
134
- height: int,
135
- width: int,
136
- num_frames: int,
137
- frame_rate: float,
138
- num_inference_steps: int,
139
- video_guider_params: MultiModalGuiderParams,
140
- audio_guider_params: MultiModalGuiderParams,
141
- images: list,
142
- tiling_config: TilingConfig | None = None,
143
- enhance_prompt: bool = False,
144
- ):
145
- """
146
- Generate video with cached LoRA state and CFG guidance.
147
- """
148
- assert_resolution(height=height, width=width, is_two_stage=True)
149
-
150
- device = self.device
151
- dtype = self.dtype
152
- generator = torch.Generator(device=device).manual_seed(seed)
153
- noiser = GaussianNoiser(generator=generator)
154
-
155
- # Apply cached LoRA state if available
156
- if self._cached_state_stage1 is not None:
157
- print("[LoRA] Applying cached state to stage 1 transformer...")
158
- t1 = self.stage_1_model_ledger.transformer()
159
- with torch.no_grad():
160
- t1.load_state_dict(self._cached_state_stage1, strict=False)
161
-
162
- if self._cached_state_stage2 is not None:
163
- print("[LoRA] Applying cached state to stage 2 transformer...")
164
- t2 = self.stage_2_model_ledger.transformer()
165
- with torch.no_grad():
166
- t2.load_state_dict(self._cached_state_stage2, strict=False)
167
-
168
- # Encode prompts
169
- from ltx_pipelines.utils.helpers import encode_prompts
170
- ctx_p, ctx_n = encode_prompts(
171
- [prompt, negative_prompt],
172
- self.stage_1_model_ledger,
173
- enhance_first_prompt=enhance_prompt,
174
- enhance_prompt_image=images[0][0] if len(images) > 0 else None,
175
- enhance_prompt_seed=seed,
176
- )
177
-
178
- v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
179
- v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
180
-
181
- # Stage 1
182
- stage_1_output_shape = VideoPixelShape(
183
- batch=1, frames=num_frames,
184
- width=width // 2, height=height // 2, fps=frame_rate
185
- )
186
-
187
- video_encoder = self.stage_1_model_ledger.video_encoder()
188
- stage_1_conditionings = combined_image_conditionings(
189
- images=images,
190
- height=stage_1_output_shape.height,
191
- width=stage_1_output_shape.width,
192
- video_encoder=video_encoder,
193
- dtype=dtype,
194
- device=device,
195
- )
196
- torch.cuda.synchronize()
197
- del video_encoder
198
- cleanup_memory()
199
-
200
- transformer = self.stage_1_model_ledger.transformer()
201
-
202
- empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
203
- stepper = Res2sDiffusionStep()
204
- sigmas = (
205
- LTX2Scheduler()
206
- .execute(latent=empty_latent, steps=num_inference_steps)
207
- .to(dtype=torch.float32, device=device)
208
- )
209
-
210
- # Stage 1 denoising with CFG
211
- from ltx_pipelines.utils.helpers import multi_modal_guider_denoising_func, res2s_audio_video_denoising_loop
212
-
213
- def first_stage_denoising_loop(sigmas, video_state, audio_state, stepper):
214
- return res2s_audio_video_denoising_loop(
215
- sigmas=sigmas,
216
- video_state=video_state,
217
- audio_state=audio_state,
218
- stepper=stepper,
219
- denoise_fn=multi_modal_guider_denoising_func(
220
- video_guider=MultiModalGuider(params=video_guider_params, negative_context=v_context_n),
221
- audio_guider=MultiModalGuider(params=audio_guider_params, negative_context=a_context_n),
222
- v_context=v_context_p,
223
- a_context=a_context_p,
224
- transformer=transformer,
225
- ),
226
- )
227
-
228
- from ltx_pipelines.utils.helpers import denoise_audio_video
229
- video_state, audio_state = denoise_audio_video(
230
- output_shape=stage_1_output_shape,
231
- conditionings=stage_1_conditionings,
232
- noiser=noiser,
233
- sigmas=sigmas,
234
- stepper=stepper,
235
- denoising_loop_fn=first_stage_denoising_loop,
236
- components=self.pipeline_components,
237
- dtype=dtype,
238
- device=device,
239
- )
240
-
241
- torch.cuda.synchronize()
242
- del transformer
243
- cleanup_memory()
244
-
245
- # Stage 2: Upsample and refine
246
- from ltx_core.model.upsampler import upsample_video
247
- video_encoder = self.stage_1_model_ledger.video_encoder()
248
- upscaled_video_latent = upsample_video(
249
- latent=video_state.latent[:1],
250
- video_encoder=video_encoder,
251
- upsampler=self.stage_2_model_ledger.spatial_upsampler(),
252
- )
253
-
254
- stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
255
- stage_2_conditionings = combined_image_conditionings(
256
- images=images,
257
- height=stage_2_output_shape.height,
258
- width=stage_2_output_shape.width,
259
- video_encoder=video_encoder,
260
- dtype=dtype,
261
- device=device,
262
- )
263
- torch.cuda.synchronize()
264
- del video_encoder
265
- cleanup_memory()
266
-
267
- transformer = self.stage_2_model_ledger.transformer()
268
- distilled_sigmas = torch.tensor(STAGE_2_DISTILLED_SIGMA_VALUES, device=device)
269
-
270
- from ltx_pipelines.utils.helpers import simple_denoising_func
271
-
272
- def second_stage_denoising_loop(sigmas, video_state, audio_state, stepper):
273
- return res2s_audio_video_denoising_loop(
274
- sigmas=sigmas,
275
- video_state=video_state,
276
- audio_state=audio_state,
277
- stepper=stepper,
278
- denoise_fn=simple_denoising_func(
279
- video_context=v_context_p,
280
- audio_context=a_context_p,
281
- transformer=transformer,
282
- ),
283
- )
284
-
285
- video_state, audio_state = denoise_audio_video(
286
- output_shape=stage_2_output_shape,
287
- conditionings=stage_2_conditionings,
288
- noiser=noiser,
289
- sigmas=distilled_sigmas,
290
- stepper=stepper,
291
- denoising_loop_fn=second_stage_denoising_loop,
292
- components=self.pipeline_components,
293
- dtype=dtype,
294
- device=device,
295
- noise_scale=distilled_sigmas[0],
296
- initial_video_latent=upscaled_video_latent,
297
- initial_audio_latent=audio_state.latent,
298
- )
299
-
300
- torch.cuda.synchronize()
301
- del transformer
302
- cleanup_memory()
303
-
304
- # Decode
305
- from ltx_core.model.audio_vae import decode_audio as vae_decode_audio
306
- from ltx_core.model.video_vae import decode_video as vae_decode_video
307
-
308
- decoded_video = vae_decode_video(
309
- video_state.latent, self.stage_2_model_ledger.video_decoder(), tiling_config, generator
310
- )
311
- decoded_audio = vae_decode_audio(
312
- audio_state.latent, self.stage_2_model_ledger.audio_decoder(), self.stage_2_model_ledger.vocoder()
313
- )
314
-
315
- return decoded_video, decoded_audio
316
-
317
-
318
- # =============================================================================
319
- # Model Download
320
- # =============================================================================
321
-
322
- print("=" * 80)
323
- print("Downloading LTX-2.3 HQ models...")
324
- print("=" * 80)
325
-
326
- checkpoint_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-22b-dev.safetensors")
327
- spatial_upsampler_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors")
328
- distilled_lora_path = hf_hub_download(repo_id=LTX_MODEL_REPO, filename="ltx-2.3-22b-distilled-lora-384.safetensors")
329
- gemma_root = snapshot_download(repo_id=GEMMA_REPO)
330
-
331
- print(f"Dev checkpoint: {checkpoint_path}")
332
- print(f"Spatial upsampler: {spatial_upsampler_path}")
333
- print(f"Distilled LoRA: {distilled_lora_path}")
334
- print(f"Gemma root: {gemma_root}")
335
-
336
- # =============================================================================
337
- # Download Custom LoRAs (from your existing app.py)
338
- # =============================================================================
339
-
340
- LORA_REPO = "dagloop5/LoRA"
341
-
342
- print("=" * 80)
343
- print("Downloading custom LoRA adapters...")
344
- print("=" * 80)
345
-
346
- pose_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="LTX2_3_NSFW_furry_concat_v2.safetensors")
347
- general_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="LTX2.3_reasoning_I2V_V3.safetensors")
348
- motion_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="motion_helper.safetensors")
349
- dreamlay_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="DR34ML4Y_LTXXX_PREVIEW_RC1.safetensors")
350
- mself_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="Furry Hyper Masturbation - LTX-2 I2V v1.safetensors")
351
- dramatic_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="LTX-2.3 - Orgasm.safetensors")
352
- fluid_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="cr3ampi3_animation_i2v_ltx2_v1.0.safetensors")
353
- liquid_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="liquid_wet_dr1pp_ltx2_v1.0_scaled.safetensors")
354
- demopose_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="clapping-cheeks-audio-v001-alpha.safetensors")
355
- voice_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="hentai_voice_ltx23.safetensors")
356
- realism_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="FurryenhancerLTX2.3V1.215.safetensors")
357
- transition_lora_path = hf_hub_download(repo_id=LORA_REPO, filename="LTX-2_takerpov_lora_v1.2.safetensors")
358
-
359
- print(f"All {12} custom LoRAs downloaded")
360
- print("=" * 80)
361
-
362
- # =============================================================================
363
- # Pipeline Initialization (with empty loras - LoRAs applied via cache)
364
- # =============================================================================
365
-
366
- print("Initializing HQ Pipeline with LoRA cache support...")
367
-
368
- distilled_lora = [
369
- LoraPathStrengthAndSDOps(
370
- path=distilled_lora_path,
371
- strength=1.0,
372
- sd_ops=LTXV_LORA_COMFY_RENAMING_MAP,
373
- )
374
- ]
375
-
376
- pipeline = HQPipelineWithCachedLoRA(
377
- checkpoint_path=checkpoint_path,
378
- distilled_lora=distilled_lora,
379
- distilled_lora_strength_stage_1=0.25,
380
- distilled_lora_strength_stage_2=0.50,
381
- spatial_upsampler_path=spatial_upsampler_path,
382
- gemma_root=gemma_root,
383
- loras=(), # LoRAs will be applied via cached state
384
- quantization=QuantizationPolicy.fp8_cast(),
385
- )
386
-
387
- print("Pipeline initialized!")
388
- print("=" * 80)
389
-
390
- # =============================================================================
391
- # ZeroGPU Tensor Preloading
392
- # =============================================================================
393
-
394
- print("Preloading all models for ZeroGPU tensor packing...")
395
-
396
- # Preload stage 1 models
397
- ledger1 = pipeline.stage_1_model_ledger
398
- _transformer_s1 = ledger1.transformer()
399
- ledger1.transformer = lambda: _transformer_s1
400
-
401
- _video_encoder_s1 = ledger1.video_encoder()
402
- ledger1.video_encoder = lambda: _video_encoder_s1
403
-
404
- _video_decoder_s1 = ledger1.video_decoder()
405
- ledger1.video_decoder = lambda: _video_decoder_s1
406
-
407
- _audio_decoder_s1 = ledger1.audio_decoder()
408
- ledger1.audio_decoder = lambda: _audio_decoder_s1
409
-
410
- _vocoder_s1 = ledger1.vocoder()
411
- ledger1.vocoder = lambda: _vocoder_s1
412
-
413
- _spatial_upsampler_s1 = ledger1.spatial_upsampler()
414
- ledger1.spatial_upsampler = lambda: _spatial_upsampler_s1
415
-
416
- _text_encoder_s1 = ledger1.text_encoder()
417
- ledger1.text_encoder = lambda: _text_encoder_s1
418
-
419
- _embeddings_processor_s1 = ledger1.embeddings_processor()
420
- ledger1.embeddings_processor = lambda: _embeddings_processor_s1
421
-
422
- print(" Stage 1 models preloaded")
423
-
424
- # Preload stage 2 models
425
- ledger2 = pipeline.stage_2_model_ledger
426
- _transformer_s2 = ledger2.transformer()
427
- ledger2.transformer = lambda: _transformer_s2
428
-
429
- _video_encoder_s2 = ledger2.video_encoder()
430
- ledger2.video_encoder = lambda: _video_encoder_s2
431
-
432
- _video_decoder_s2 = ledger2.video_decoder()
433
- ledger2.video_decoder = lambda: _video_decoder_s2
434
-
435
- _audio_decoder_s2 = ledger2.audio_decoder()
436
- ledger2.audio_decoder = lambda: _audio_decoder_s2
437
-
438
- _vocoder_s2 = ledger2.vocoder()
439
- ledger2.vocoder = lambda: _vocoder_s2
440
-
441
- _spatial_upsampler_s2 = ledger2.spatial_upsampler()
442
- ledger2.spatial_upsampler = lambda: _spatial_upsampler_s2
443
-
444
- _text_encoder_s2 = ledger2.text_encoder()
445
- ledger2.text_encoder = lambda: _text_encoder_s2
446
-
447
- _embeddings_processor_s2 = ledger2.embeddings_processor()
448
- ledger2.embeddings_processor = lambda: _embeddings_processor_s2
449
-
450
- print(" Stage 2 models preloaded")
451
-
452
- print("All models preloaded for ZeroGPU tensor packing!")
453
- print("=" * 80)
454
-
455
- # =============================================================================
456
- # LoRA Cache Functions (from your existing app.py)
457
- # =============================================================================
458
-
459
- LORA_CACHE_DIR = Path("lora_cache")
460
- LORA_CACHE_DIR.mkdir(exist_ok=True)
461
-
462
- def prepare_lora_cache(
463
- pose_strength: float, general_strength: float, motion_strength: float,
464
- dreamlay_strength: float, mself_strength: float, dramatic_strength: float,
465
- fluid_strength: float, liquid_strength: float, demopose_strength: float,
466
- voice_strength: float, realism_strength: float, transition_strength: float,
467
- progress=gr.Progress(track_tqdm=True),
468
- ):
469
- """
470
- Build cached LoRA state for both stages (different strengths).
471
- """
472
- global _pipeline
473
-
474
- progress(0.05, desc="Preparing LoRA cache...")
475
-
476
- # Create key for cache
477
- key_str = f"{checkpoint_path}:{pose_strength}:{general_strength}:{motion_strength}:{dreamlay_strength}:{mself_strength}:{dramatic_strength}:{fluid_strength}:{liquid_strength}:{demopose_strength}:{voice_strength}:{realism_strength}:{transition_strength}"
478
- key = hashlib.sha256(key_str.encode()).hexdigest()
479
-
480
- cache_path_stage1 = LORA_CACHE_DIR / f"{key}_stage1.safetensors"
481
- cache_path_stage2 = LORA_CACHE_DIR / f"{key}_stage2.safetensors"
482
-
483
- # Check if cached
484
- if cache_path_stage1.exists() and cache_path_stage2.exists():
485
- progress(0.20, desc="Loading cached LoRA state...")
486
- state_stage1 = load_file(str(cache_path_stage1))
487
- state_stage2 = load_file(str(cache_path_stage2))
488
- pipeline.apply_cached_lora_state(state_stage1, state_stage2)
489
- return f"Loaded cached LoRA state: {cache_path_stage1.name}"
490
-
491
- # Build LoRA list (distilled + 12 custom)
492
- entries = [
493
- (distilled_lora_path, 0.25), # Stage 1 strength
494
- (pose_lora_path, pose_strength),
495
- (general_lora_path, general_strength),
496
- (motion_lora_path, motion_strength),
497
- (dreamlay_lora_path, dreamlay_strength),
498
- (mself_lora_path, mself_strength),
499
- (dramatic_lora_path, dramatic_strength),
500
- (fluid_lora_path, fluid_strength),
501
- (liquid_lora_path, liquid_strength),
502
- (demopose_lora_path, demopose_strength),
503
- (voice_lora_path, voice_strength),
504
- (realism_lora_path, realism_strength),
505
- (transition_lora_path, transition_strength),
506
- ]
507
-
508
- loras = [
509
- LoraPathStrengthAndSDOps(path, strength, LTXV_LORA_COMFY_RENAMING_MAP)
510
- for path, strength in entries
511
- if path is not None and float(strength) != 0.0
512
- ]
513
-
514
- # Build stage 1 cached state
515
- progress(0.35, desc="Building stage 1 fused state (CPU)...")
516
- tmp_ledger1 = pipeline.stage_1_model_ledger.__class__(
517
- dtype=torch.bfloat16,
518
- device=torch.device("cpu"),
519
- checkpoint_path=str(checkpoint_path),
520
- spatial_upsampler_path=str(spatial_upsampler_path),
521
- gemma_root_path=str(gemma_root),
522
- loras=tuple(loras),
523
- quantization=None, # No quantization for CPU cache build
524
- )
525
- transformer1 = tmp_ledger1.transformer()
526
- state_stage1 = {k: v.detach().cpu().contiguous() for k, v in transformer1.state_dict().items()}
527
- save_file(state_stage1, str(cache_path_stage1))
528
-
529
- del transformer1, tmp_ledger1
530
- gc.collect()
531
-
532
- # Build stage 2 cached state (different LoRA strengths)
533
- progress(0.65, desc="Building stage 2 fused state (CPU)...")
534
-
535
- # Update strengths for stage 2
536
- entries_s2 = [
537
- (distilled_lora_path, 0.50), # Stage 2 strength
538
- (pose_lora_path, pose_strength),
539
- (general_lora_path, general_strength),
540
- (motion_lora_path, motion_strength),
541
- (dreamlay_lora_path, dreamlay_strength),
542
- (mself_lora_path, mself_strength),
543
- (dramatic_lora_path, dramatic_strength),
544
- (fluid_lora_path, fluid_strength),
545
- (liquid_lora_path, liquid_strength),
546
- (demopose_lora_path, demopose_strength),
547
- (voice_lora_path, voice_strength),
548
- (realism_lora_path, realism_strength),
549
- (transition_lora_path, transition_strength),
550
- ]
551
-
552
- loras_s2 = [
553
- LoraPathStrengthAndSDOps(path, strength, LTXV_LORA_COMFY_RENAMING_MAP)
554
- for path, strength in entries_s2
555
- if path is not None and float(strength) != 0.0
556
- ]
557
-
558
- tmp_ledger2 = pipeline.stage_2_model_ledger.__class__(
559
- dtype=torch.bfloat16,
560
- device=torch.device("cpu"),
561
- checkpoint_path=str(checkpoint_path),
562
- spatial_upsampler_path=str(spatial_upsampler_path),
563
- gemma_root_path=str(gemma_root),
564
- loras=tuple(loras_s2),
565
- quantization=None,
566
- )
567
- transformer2 = tmp_ledger2.transformer()
568
- state_stage2 = {k: v.detach().cpu().contiguous() for k, v in transformer2.state_dict().items()}
569
- save_file(state_stage2, str(cache_path_stage2))
570
-
571
- del transformer2, tmp_ledger2
572
- gc.collect()
573
-
574
- progress(0.90, desc="Applying LoRA state to pipeline...")
575
- pipeline.apply_cached_lora_state(state_stage1, state_stage2)
576
-
577
- progress(1.0, desc="Done!")
578
- return f"Built and cached LoRA state: {cache_path_stage1.name}"
579
-
580
-
581
- # =============================================================================
582
- # Helper Functions
583
- # =============================================================================
584
-
585
- def log_memory(tag: str):
586
- if torch.cuda.is_available():
587
- allocated = torch.cuda.memory_allocated() / 1024**3
588
- peak = torch.cuda.max_memory_allocated() / 1024**3
589
- free, total = torch.cuda.mem_get_info()
590
- print(f"[VRAM {tag}] allocated={allocated:.2f}GB peak={peak:.2f}GB free={free / 1024**3:.2f}GB total={total / 1024**3:.2f}GB")
591
-
592
-
593
- def calculate_frames(duration: float, frame_rate: float = DEFAULT_FRAME_RATE) -> int:
594
- ideal_frames = int(duration * frame_rate)
595
- ideal_frames = max(ideal_frames, MIN_FRAMES)
596
- k = round((ideal_frames - 1) / 8)
597
- frames = k * 8 + 1
598
- return min(frames, MAX_FRAMES)
599
-
600
-
601
- def validate_resolution(height: int, width: int) -> tuple[int, int]:
602
- height = round(height / STEP) * STEP
603
- width = round(width / STEP) * STEP
604
- height = max(MIN_DIM, min(height, MAX_DIM))
605
- width = max(MIN_DIM, min(width, MAX_DIM))
606
- return height, width
607
-
608
-
609
- def detect_aspect_ratio(image) -> str:
610
- if image is None:
611
- return "16:9"
612
- if hasattr(image, "size"):
613
- w, h = image.size
614
- elif hasattr(image, "shape"):
615
- h, w = image.shape[:2]
616
- else:
617
- return "16:9"
618
- ratio = w / h
619
- candidates = {"16:9": 16/9, "9:16": 9/16, "1:1": 1.0}
620
- return min(candidates, key=lambda k: abs(ratio - candidates[k]))
621
-
622
-
623
- def get_duration(
624
- prompt, negative_prompt, first_image, last_image, input_audio,
625
- duration, seed, randomize_seed, height, width, enhance_prompt,
626
- video_cfg_scale, video_stg_scale, video_rescale_scale, video_a2v_scale,
627
- audio_cfg_scale, audio_stg_scale, audio_rescale_scale, audio_v2a_scale,
628
- pose_strength, general_strength, motion_strength, dreamlay_strength,
629
- mself_strength, dramatic_strength, fluid_strength, liquid_strength,
630
- demopose_strength, voice_strength, realism_strength, transition_strength,
631
- progress=None,
632
- ) -> int:
633
- base = 90
634
- if duration > 4:
635
- base += 15
636
- if duration > 6:
637
- base += 15
638
- if height > 700 or width > 1000:
639
- base += 15
640
- return min(base, 120)
641
-
642
-
643
- @spaces.GPU(duration=get_duration)
644
- @torch.inference_mode()
645
- def generate_video(
646
- prompt: str,
647
- negative_prompt: str,
648
- first_image,
649
- last_image,
650
- input_audio,
651
- duration: float,
652
- seed: int,
653
- randomize_seed: bool,
654
- height: int,
655
- width: int,
656
- enhance_prompt: bool,
657
- video_cfg_scale: float,
658
- video_stg_scale: float,
659
- video_rescale_scale: float,
660
- video_a2v_scale: float,
661
- audio_cfg_scale: float,
662
- audio_stg_scale: float,
663
- audio_rescale_scale: float,
664
- audio_v2a_scale: float,
665
- pose_strength: float,
666
- general_strength: float,
667
- motion_strength: float,
668
- dreamlay_strength: float,
669
- mself_strength: float,
670
- dramatic_strength: float,
671
- fluid_strength: float,
672
- liquid_strength: float,
673
- demopose_strength: float,
674
- voice_strength: float,
675
- realism_strength: float,
676
- transition_strength: float,
677
- progress=gr.Progress(track_tqdm=True),
678
- ):
679
- try:
680
- torch.cuda.reset_peak_memory_stats()
681
- log_memory("start")
682
-
683
- current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)
684
- print(f"Using seed: {current_seed}")
685
-
686
- height, width = validate_resolution(int(height), int(width))
687
- print(f"Resolution: {width}x{height}")
688
-
689
- num_frames = calculate_frames(duration, DEFAULT_FRAME_RATE)
690
- print(f"Frames: {num_frames} ({duration}s @ {DEFAULT_FRAME_RATE}fps)")
691
-
692
- images = []
693
- output_dir = Path("outputs")
694
- output_dir.mkdir(exist_ok=True)
695
-
696
- if first_image is not None:
697
- temp_first_path = output_dir / f"temp_first_{current_seed}.jpg"
698
- if hasattr(first_image, "save"):
699
- first_image.save(temp_first_path)
700
- else:
701
- import shutil
702
- shutil.copy(first_image, temp_first_path)
703
- images.append((str(temp_first_path), 1.0))
704
-
705
- if last_image is not None:
706
- temp_last_path = output_dir / f"temp_last_{current_seed}.jpg"
707
- if hasattr(last_image, "save"):
708
- last_image.save(temp_last_path)
709
- else:
710
- import shutil
711
- shutil.copy(last_image, temp_last_path)
712
- images.append((str(temp_last_path), 1.0))
713
-
714
- tiling_config = TilingConfig.default()
715
- video_chunks_number = get_video_chunks_number(num_frames, tiling_config)
716
-
717
- video_guider_params = MultiModalGuiderParams(
718
- cfg_scale=video_cfg_scale,
719
- stg_scale=video_stg_scale,
720
- rescale_scale=video_rescale_scale,
721
- modality_scale=video_a2v_scale,
722
- skip_step=0,
723
- stg_blocks=[],
724
- )
725
-
726
- audio_guider_params = MultiModalGuiderParams(
727
- cfg_scale=audio_cfg_scale,
728
- stg_scale=audio_stg_scale,
729
- rescale_scale=audio_rescale_scale,
730
- modality_scale=audio_v2a_scale,
731
- skip_step=0,
732
- stg_blocks=[],
733
- )
734
-
735
- log_memory("before pipeline call")
736
-
737
- video, audio = pipeline(
738
- prompt=prompt,
739
- negative_prompt=negative_prompt,
740
- seed=current_seed,
741
- height=height,
742
- width=width,
743
- num_frames=num_frames,
744
- frame_rate=DEFAULT_FRAME_RATE,
745
- num_inference_steps=LTX_2_3_HQ_PARAMS.num_inference_steps,
746
- video_guider_params=video_guider_params,
747
- audio_guider_params=audio_guider_params,
748
- images=images,
749
- tiling_config=tiling_config,
750
- enhance_prompt=enhance_prompt,
751
- )
752
-
753
- log_memory("after pipeline call")
754
-
755
- output_path = tempfile.mktemp(suffix=".mp4")
756
- encode_video(
757
- video=video,
758
- fps=DEFAULT_FRAME_RATE,
759
- audio=audio,
760
- output_path=output_path,
761
- video_chunks_number=video_chunks_number,
762
- )
763
-
764
- log_memory("after encode_video")
765
- return str(output_path), current_seed
766
-
767
- except Exception as e:
768
- import traceback
769
- log_memory("on error")
770
- print(f"Error: {str(e)}\n{traceback.format_exc()}")
771
- return None, current_seed
772
-
773
-
774
- # =============================================================================
775
- # Gradio UI (Merged from your app.py)
776
- # =============================================================================
777
-
778
- css = """
779
- .fillable {max-width: 1200px !important}
780
- .progress-text {color: white}
781
- """
782
-
783
- with gr.Blocks(title="LTX-2.3 Two-Stage HQ with LoRA Cache") as demo:
784
- gr.Markdown("# LTX-2.3 Two-Stage HQ Video Generation with LoRA Cache")
785
- gr.Markdown(
786
- "High-quality text/image-to-video with cached LoRA state + CFG guidance. "
787
- "[[Model]](https://huggingface.co/Lightricks/LTX-2.3)"
788
- )
789
-
790
- with gr.Row():
791
- with gr.Column():
792
- with gr.Row():
793
- first_image = gr.Image(label="First Frame (Optional)", type="pil")
794
- last_image = gr.Image(label="Last Frame (Optional)", type="pil")
795
-
796
- prompt = gr.Textbox(
797
- label="Prompt",
798
- value=DEFAULT_PROMPT,
799
- lines=3,
800
- )
801
-
802
- negative_prompt = gr.Textbox(
803
- label="Negative Prompt",
804
- value=DEFAULT_NEGATIVE_PROMPT,
805
- lines=2,
806
- )
807
-
808
- duration = gr.Slider(
809
- label="Duration (seconds)",
810
- minimum=0.5, maximum=8.0, value=2.0, step=0.1,
811
- )
812
-
813
- enhance_prompt = gr.Checkbox(label="Enhance Prompt", value=False)
814
-
815
- generate_btn = gr.Button("Generate Video", variant="primary", size="lg")
816
-
817
- with gr.Column():
818
- output_video = gr.Video(label="Generated Video", autoplay=True)
819
-
820
- with gr.Accordion("Advanced Settings", open=False):
821
- with gr.Row():
822
- width = gr.Number(label="Width", value=1280, precision=0)
823
- height = gr.Number(label="Height", value=704, precision=0)
824
-
825
- with gr.Row():
826
- seed = gr.Number(label="Seed", value=42, precision=0, minimum=0, maximum=MAX_SEED)
827
- randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
828
-
829
- gr.Markdown("### Video Guidance Parameters")
830
-
831
- with gr.Row():
832
- video_cfg_scale = gr.Slider(
833
- label="Video CFG Scale", minimum=1.0, maximum=10.0,
834
- value=LTX_2_3_HQ_PARAMS.video_guider_params.cfg_scale, step=0.1
835
- )
836
- video_stg_scale = gr.Slider(
837
- label="Video STG Scale", minimum=0.0, maximum=2.0, value=0.0, step=0.1
838
- )
839
-
840
- with gr.Row():
841
- video_rescale_scale = gr.Slider(
842
- label="Video Rescale", minimum=0.0, maximum=2.0, value=0.45, step=0.1
843
- )
844
- video_a2v_scale = gr.Slider(
845
- label="A2V Scale", minimum=0.0, maximum=5.0, value=3.0, step=0.1
846
- )
847
-
848
- gr.Markdown("### Audio Guidance Parameters")
849
-
850
- with gr.Row():
851
- audio_cfg_scale = gr.Slider(
852
- label="Audio CFG Scale", minimum=1.0, maximum=15.0,
853
- value=LTX_2_3_HQ_PARAMS.audio_guider_params.cfg_scale, step=0.1
854
- )
855
- audio_stg_scale = gr.Slider(
856
- label="Audio STG Scale", minimum=0.0, maximum=2.0, value=0.0, step=0.1
857
- )
858
-
859
- with gr.Row():
860
- audio_rescale_scale = gr.Slider(
861
- label="Audio Rescale", minimum=0.0, maximum=2.0, value=1.0, step=0.1
862
- )
863
- audio_v2a_scale = gr.Slider(
864
- label="V2A Scale", minimum=0.0, maximum=5.0, value=3.0, step=0.1
865
- )
866
-
867
- gr.Markdown("### LoRA Adapter Strengths")
868
-
869
- with gr.Row():
870
- pose_strength = gr.Slider(label="Anthro Enhancer", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
871
- general_strength = gr.Slider(label="Reasoning Enhancer", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
872
-
873
- with gr.Row():
874
- motion_strength = gr.Slider(label="Anthro Posing", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
875
- dreamlay_strength = gr.Slider(label="Dreamlay", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
876
-
877
- with gr.Row():
878
- mself_strength = gr.Slider(label="Mself", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
879
- dramatic_strength = gr.Slider(label="Dramatic", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
880
-
881
- with gr.Row():
882
- fluid_strength = gr.Slider(label="Fluid Helper", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
883
- liquid_strength = gr.Slider(label="Liquid Helper", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
884
-
885
- with gr.Row():
886
- demopose_strength = gr.Slider(label="Audio Helper", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
887
- voice_strength = gr.Slider(label="Voice Helper", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
888
-
889
- with gr.Row():
890
- realism_strength = gr.Slider(label="Anthro Realism", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
891
- transition_strength = gr.Slider(label="POV", minimum=0.0, maximum=2.0, value=0.0, step=0.01)
892
-
893
- prepare_lora_btn = gr.Button("Prepare / Load LoRA Cache", variant="secondary")
894
- lora_status = gr.Textbox(
895
- label="LoRA Cache Status",
896
- value="No LoRA state prepared yet.",
897
- interactive=False,
898
- )
899
-
900
- def on_image_upload(image, current_h, current_w):
901
- if image is None:
902
- return gr.update(), gr.update()
903
- aspect = detect_aspect_ratio(image)
904
- if aspect in RESOLUTIONS:
905
- return (
906
- gr.update(value=RESOLUTIONS[aspect]["width"]),
907
- gr.update(value=RESOLUTIONS[aspect]["height"])
908
- )
909
- return gr.update(), gr.update()
910
-
911
- first_image.change(fn=on_image_upload, inputs=[first_image, height, width], outputs=[width, height])
912
- last_image.change(fn=on_image_upload, inputs=[last_image, height, width], outputs=[width, height])
913
-
914
- prepare_lora_btn.click(
915
- fn=prepare_lora_cache,
916
- inputs=[pose_strength, general_strength, motion_strength, dreamlay_strength,
917
- mself_strength, dramatic_strength, fluid_strength, liquid_strength,
918
- demopose_strength, voice_strength, realism_strength, transition_strength],
919
- outputs=[lora_status],
920
- )
921
-
922
- generate_btn.click(
923
- fn=generate_video,
924
- inputs=[
925
- prompt, negative_prompt, first_image, last_image, None, duration,
926
- seed, randomize_seed, height, width, enhance_prompt,
927
- video_cfg_scale, video_stg_scale, video_rescale_scale, video_a2v_scale,
928
- audio_cfg_scale, audio_stg_scale, audio_rescale_scale, audio_v2a_scale,
929
- pose_strength, general_strength, motion_strength, dreamlay_strength,
930
- mself_strength, dramatic_strength, fluid_strength, liquid_strength,
931
- demopose_strength, voice_strength, realism_strength, transition_strength,
932
- ],
933
- outputs=[output_video, seed],
934
- )
935
-
936
-
937
- if __name__ == "__main__":
938
- demo.queue().launch(theme=gr.themes.Citrus(), css=css, mcp_server=True)