Create app(draft).py

app(draft).py

@@ -0,0 +1,373 @@
+import os
+import subprocess
+import sys
+
+# Disable torch.compile / dynamo before any torch import
+os.environ["TORCH_COMPILE_DISABLE"] = "1"
+os.environ["TORCHDYNAMO_DISABLE"] = "1"
+
+# Install xformers for memory-efficient attention
+subprocess.run([sys.executable, "-m", "pip", "install", "xformers==0.0.32.post2", "--no-build-isolation"], check=False)
+
+# Clone LTX-2 repo and install packages
+LTX_REPO_URL = "https://github.com/Lightricks/LTX-2.git"
+LTX_REPO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "LTX-2")
+LTX_COMMIT_SHA = "a2c3f24078eb918171967f74b6f66b756b29ee45"
+
+if not os.path.exists(LTX_REPO_DIR):
+    print(f"Cloning {LTX_REPO_URL}...")
+    os.makedirs(LTX_REPO_DIR)
+    subprocess.run(["git", "init", LTX_REPO_DIR], check=True)
+    subprocess.run(["git", "remote", "add", "origin", LTX_REPO_URL], cwd=LTX_REPO_DIR, check=True)
+    subprocess.run(["git", "fetch", "--depth", "1", "origin", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True)
+    subprocess.run(["git", "checkout", LTX_COMMIT_SHA], cwd=LTX_REPO_DIR, check=True)
+
+
+print("Installing ltx-core and ltx-pipelines from cloned repo...")
+subprocess.run(
+    [sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "-e",
+     os.path.join(LTX_REPO_DIR, "packages", "ltx-core"),
+     "-e", os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines")],
+    check=True,
+)
+
+sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-pipelines", "src"))
+sys.path.insert(0, os.path.join(LTX_REPO_DIR, "packages", "ltx-core", "src"))
+
+import logging
+import random
+import tempfile
+from pathlib import Path
+from collections.abc import Iterator
+
+import torch
+torch._dynamo.config.suppress_errors = True
+torch._dynamo.config.disable = True
+
+import spaces
+import gradio as gr
+import numpy as np
+from huggingface_hub import hf_hub_download, snapshot_download
+
+from ltx_core.components.diffusion_steps import Res2sDiffusionStep
+from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams
+from ltx_core.components.noisers import GaussianNoiser
+from ltx_core.components.schedulers import LTX2Scheduler
+from ltx_core.loader import LoraPathStrengthAndSDOps
+from ltx_core.loader.registry import Registry
+from ltx_core.model.video_vae import TilingConfig, get_video_chunks_number
+from ltx_core.quantization import QuantizationPolicy
+from ltx_core.types import Audio, VideoLatentShape, VideoPixelShape
+from ltx_pipelines.utils.args import ImageConditioningInput, hq_2_stage_arg_parser
+from ltx_pipelines.utils.blocks import (
+    AudioDecoder,
+    DiffusionStage,
+    ImageConditioner,
+    PromptEncoder,
+    VideoDecoder,
+    VideoUpsampler,
+)
+from ltx_pipelines.utils.constants import (
+    LTX_2_3_HQ_PARAMS,
+    STAGE_2_DISTILLED_SIGMAS,
+)
+from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser
+from ltx_pipelines.utils.helpers import (
+    assert_resolution,
+    combined_image_conditionings,
+    get_device,
+)
+from ltx_pipelines.utils.media_io import encode_video
+from ltx_pipelines.utils.samplers import res2s_audio_video_denoising_loop
+from ltx_pipelines.utils.types import ModalitySpec
+
+# Force-patch xformers attention into the LTX attention module.
+from ltx_core.model.transformer import attention as _attn_mod
+print(f"[ATTN] Before patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
+try:
+    from xformers.ops import memory_efficient_attention as _mea
+    _attn_mod.memory_efficient_attention = _mea
+    print(f"[ATTN] After patch: memory_efficient_attention={_attn_mod.memory_efficient_attention}")
+except Exception as e:
+    print(f"[ATTN] xformers patch FAILED: {type(e).__name__}: {e}")
+
+logging.getLogger().setLevel(logging.INFO)
+
+MAX_SEED = np.iinfo(np.int32).max
+DEFAULT_PROMPT = (
+    "An astronaut hatches from a fragile egg on the surface of the Moon, "
+    "the shell cracking and peeling apart in gentle low-gravity motion. "
+    "Fine lunar dust lifts and drifts outward with each movement, floating "
+    "in slow arcs before settling back onto the ground."
+)
+DEFAULT_FRAME_RATE = 24.0
+
+# Resolution presets: (width, height)
+RESOLUTIONS = {
+    "high": {"16:9": (1536, 1024), "9:16": (1024, 1536), "1:1": (1024, 1024)},
+    "low": {"16:9": (768, 512), "9:16": (512, 768), "1:1": (768, 768)},
+}
+
+class TI2VidTwoStagesHQPipeline:
+    """
+    Two-stage text/image-to-video generation pipeline using the res_2s sampler.
+    Same structure as :class:`TI2VidTwoStagesPipeline`: stage 1 generates video at
+    half of the target resolution with CFG guidance (assuming  full model is used),
+    then Stage 2 upsamples by 2x and refines using a distilled LoRA for higher
+    quality output.
+    Uses the res_2s second-order sampler instead of Euler, allowing fewer
+    steps for comparable quality. Supports optional image conditioning via
+    the images parameter.
+    """
+
+    def __init__(  # noqa: PLR0913
+        self,
+        checkpoint_path: str,
+        distilled_lora: list[LoraPathStrengthAndSDOps],
+        distilled_lora_strength_stage_1: float,
+        distilled_lora_strength_stage_2: float,
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        loras: tuple[LoraPathStrengthAndSDOps, ...],
+        device: torch.device | None = None,
+        quantization: QuantizationPolicy | None = None,
+        registry: Registry | None = None,
+        torch_compile: bool = False,
+    ):
+        self.device = device or get_device()
+        self.dtype = torch.bfloat16
+        self._scheduler = LTX2Scheduler()
+
+        distilled_lora_stage_1 = LoraPathStrengthAndSDOps(
+            path=distilled_lora[0].path,
+            strength=distilled_lora_strength_stage_1,
+            sd_ops=distilled_lora[0].sd_ops,
+        )
+        distilled_lora_stage_2 = LoraPathStrengthAndSDOps(
+            path=distilled_lora[0].path,
+            strength=distilled_lora_strength_stage_2,
+            sd_ops=distilled_lora[0].sd_ops,
+        )
+
+        self.prompt_encoder = PromptEncoder(checkpoint_path, gemma_root, self.dtype, self.device, registry=registry)
+        self.image_conditioner = ImageConditioner(checkpoint_path, self.dtype, self.device, registry=registry)
+        self.upsampler = VideoUpsampler(
+            checkpoint_path, spatial_upsampler_path, self.dtype, self.device, registry=registry
+        )
+        self.video_decoder = VideoDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
+        self.audio_decoder = AudioDecoder(checkpoint_path, self.dtype, self.device, registry=registry)
+
+        self.stage_1 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=(*loras, distilled_lora_stage_1),
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
+        )
+        self.stage_2 = DiffusionStage(
+            checkpoint_path,
+            self.dtype,
+            self.device,
+            loras=(*loras, distilled_lora_stage_2),
+            quantization=quantization,
+            registry=registry,
+            torch_compile=torch_compile,
+        )
+
+    @torch.inference_mode()
+    def __call__(  # noqa: PLR0913
+        self,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        height: int,
+        width: int,
+        num_frames: int,
+        frame_rate: float,
+        num_inference_steps: int,
+        video_guider_params: MultiModalGuiderParams,
+        audio_guider_params: MultiModalGuiderParams,
+        images: list[ImageConditioningInput],
+        tiling_config: TilingConfig | None = None,
+        enhance_prompt: bool = False,
+        streaming_prefetch_count: int | None = None,
+        max_batch_size: int = 1,
+        stage_1_sigmas: torch.Tensor | None = None,
+        stage_2_sigmas: torch.Tensor = STAGE_2_DISTILLED_SIGMAS,
+    ) -> tuple[Iterator[torch.Tensor], Audio]:
+        assert_resolution(height=height, width=width, is_two_stage=True)
+
+        generator = torch.Generator(device=self.device).manual_seed(seed)
+        noiser = GaussianNoiser(generator=generator)
+        dtype = torch.bfloat16
+
+        ctx_p, ctx_n = self.prompt_encoder(
+            [prompt, negative_prompt],
+            enhance_first_prompt=enhance_prompt,
+            enhance_prompt_image=images[0][0] if len(images) > 0 else None,
+            enhance_prompt_seed=seed,
+            streaming_prefetch_count=streaming_prefetch_count,
+        )
+        v_context_p, a_context_p = ctx_p.video_encoding, ctx_p.audio_encoding
+        v_context_n, a_context_n = ctx_n.video_encoding, ctx_n.audio_encoding
+
+        # Stage 1: Generate video at half resolution with CFG guidance using res2s sampler.
+        stage_1_output_shape = VideoPixelShape(
+            batch=1,
+            frames=num_frames,
+            width=width // 2,
+            height=height // 2,
+            fps=frame_rate,
+        )
+        stage_1_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=stage_1_output_shape.height,
+                width=stage_1_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
+        )
+
+        stepper = Res2sDiffusionStep()
+
+        if stage_1_sigmas is None:
+            empty_latent = torch.empty(VideoLatentShape.from_pixel_shape(stage_1_output_shape).to_torch_shape())
+            stage_1_sigmas = self._scheduler.execute(latent=empty_latent, steps=num_inference_steps)
+        sigmas = stage_1_sigmas.to(dtype=torch.float32, device=self.device)
+
+        video_state, audio_state = self.stage_1(
+            denoiser=GuidedDenoiser(
+                v_context=v_context_p,
+                a_context=a_context_p,
+                video_guider=MultiModalGuider(
+                    params=video_guider_params,
+                    negative_context=v_context_n,
+                ),
+                audio_guider=MultiModalGuider(
+                    params=audio_guider_params,
+                    negative_context=a_context_n,
+                ),
+            ),
+            sigmas=sigmas,
+            noiser=noiser,
+            stepper=stepper,
+            width=stage_1_output_shape.width,
+            height=stage_1_output_shape.height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(context=v_context_p, conditionings=stage_1_conditionings),
+            audio=ModalitySpec(context=a_context_p),
+            loop=res2s_audio_video_denoising_loop,
+            streaming_prefetch_count=streaming_prefetch_count,
+            max_batch_size=max_batch_size,
+        )
+
+        # Stage 2: Upsample and refine the video at higher resolution with distilled LoRA.
+        upscaled_video_latent = self.upsampler(video_state.latent[:1])
+
+        stage_2_sigmas = stage_2_sigmas.to(dtype=torch.float32, device=self.device)
+        stage_2_output_shape = VideoPixelShape(batch=1, frames=num_frames, width=width, height=height, fps=frame_rate)
+        stage_2_conditionings = self.image_conditioner(
+            lambda enc: combined_image_conditionings(
+                images=images,
+                height=stage_2_output_shape.height,
+                width=stage_2_output_shape.width,
+                video_encoder=enc,
+                dtype=dtype,
+                device=self.device,
+            )
+        )
+
+        video_state, audio_state = self.stage_2(
+            denoiser=SimpleDenoiser(v_context=v_context_p, a_context=a_context_p),
+            sigmas=stage_2_sigmas,
+            noiser=noiser,
+            stepper=stepper,
+            width=width,
+            height=height,
+            frames=num_frames,
+            fps=frame_rate,
+            video=ModalitySpec(
+                context=v_context_p,
+                conditionings=stage_2_conditionings,
+                noise_scale=stage_2_sigmas[0].item(),
+                initial_latent=upscaled_video_latent,
+            ),
+            audio=ModalitySpec(
+                context=a_context_p,
+                noise_scale=stage_2_sigmas[0].item(),
+                initial_latent=audio_state.latent,
+            ),
+            loop=res2s_audio_video_denoising_loop,
+            streaming_prefetch_count=streaming_prefetch_count,
+        )
+
+        decoded_video = self.video_decoder(video_state.latent, tiling_config, generator)
+        decoded_audio = self.audio_decoder(audio_state.latent)
+        return decoded_video, decoded_audio
+
+
+@torch.inference_mode()
+def main() -> None:
+    logging.getLogger().setLevel(logging.INFO)
+    parser = hq_2_stage_arg_parser(params=LTX_2_3_HQ_PARAMS)
+    args = parser.parse_args()
+    pipeline = TI2VidTwoStagesHQPipeline(
+        checkpoint_path=args.checkpoint_path,
+        distilled_lora=args.distilled_lora,
+        distilled_lora_strength_stage_1=args.distilled_lora_strength_stage_1,
+        distilled_lora_strength_stage_2=args.distilled_lora_strength_stage_2,
+        spatial_upsampler_path=args.spatial_upsampler_path,
+        gemma_root=args.gemma_root,
+        loras=tuple(args.lora) if args.lora else (),
+        quantization=args.quantization,
+        torch_compile=args.compile,
+    )
+    tiling_config = TilingConfig.default()
+    video_chunks_number = get_video_chunks_number(args.num_frames, tiling_config)
+    video, audio = pipeline(
+        prompt=args.prompt,
+        negative_prompt=args.negative_prompt,
+        seed=args.seed,
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        frame_rate=args.frame_rate,
+        num_inference_steps=args.num_inference_steps,
+        video_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.video_cfg_guidance_scale,
+            stg_scale=args.video_stg_guidance_scale,
+            rescale_scale=args.video_rescale_scale,
+            modality_scale=args.a2v_guidance_scale,
+            skip_step=args.video_skip_step,
+            stg_blocks=args.video_stg_blocks,
+        ),
+        audio_guider_params=MultiModalGuiderParams(
+            cfg_scale=args.audio_cfg_guidance_scale,
+            stg_scale=args.audio_stg_guidance_scale,
+            rescale_scale=args.audio_rescale_scale,
+            modality_scale=args.v2a_guidance_scale,
+            skip_step=args.audio_skip_step,
+            stg_blocks=args.audio_stg_blocks,
+        ),
+        images=args.images,
+        tiling_config=tiling_config,
+        streaming_prefetch_count=args.streaming_prefetch_count,
+        max_batch_size=args.max_batch_size,
+    )
+
+    encode_video(
+        video=video,
+        fps=args.frame_rate,
+        audio=audio,
+        output_path=args.output_path,
+        video_chunks_number=video_chunks_number,
+    )
+
+
+if __name__ == "__main__":
+    main()