app.py

import types
import random
import spaces

import torch
import numpy as np
from diffusers import LTXLatentUpsamplePipeline
from diffusers.utils import export_to_video
import gradio as gr
import tempfile

from src.transformer_ltx_nag import NAGLTXVideoTransformer3DModel
from src.pipeline_ltx_condition import NAGLTXConditionPipeline


MOD_VALUE = 32
DEFAULT_DURATION_SECONDS = 5
DEFAULT_SEED = 2025
DEFAULT_H_SLIDER_VALUE = 480
DEFAULT_W_SLIDER_VALUE = 832
NEW_FORMULA_MAX_AREA = 704.0 * 1216.0
DOWNSCALE_FACTOR = 2 / 3

SLIDER_MIN_H, SLIDER_MAX_H = 128, 1280
SLIDER_MIN_W, SLIDER_MAX_W = 128, 1280
MAX_SEED = np.iinfo(np.int32).max

FIXED_FPS = 24
MIN_FRAMES_MODEL = 8
MAX_FRAMES_MODEL = 257

DEFAULT_NAG_NEGATIVE_PROMPT = "static, motionless, still, lifeless, dull, frozen, ugly, bad quality, worst quality, poorly drawn, low resolution, blurry, lack of detail"

model_id = "Lightricks/LTX-Video-0.9.7-distilled"
transformer = NAGLTXVideoTransformer3DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
)
pipe = NAGLTXConditionPipeline.from_pretrained(
    model_id,
    transformer=transformer,
    torch_dtype=torch.bfloat16,
)
pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained("Lightricks/ltxv-spatial-upscaler-0.9.7", vae=pipe.vae, torch_dtype=torch.bfloat16)
pipe.to("cuda")
pipe_upsample.to("cuda")
pipe.vae.enable_tiling()

examples = [
    [
        "A rock star passionately plays electric guitar with intensity and emotion on a stage. The background is shrouded in deep darkness. Spotlights casts dramatic shadows.",
        DEFAULT_NAG_NEGATIVE_PROMPT,
        11,
    ],
    [
        "A clear, turquoise river flows through a rocky canyon, cascading over a small waterfall and forming a pool of water at the bottom. The river is the main focus of the scene. The overall tone of the scene is one of peace and tranquility.",
        "trees, grass, greenery",
        11,
    ],
    [
        "A woman with blood on her face and a white tank top looks down and to her right, then back up as she speaks. She has dark hair pulled back, light skin, and her face and chest are covered in blood. The camera angle is a close-up, focused on the woman's face and upper torso. The lighting is dim and blue-toned, creating a somber and intense atmosphere. The scene appears to be from a movie or TV show.",
        DEFAULT_NAG_NEGATIVE_PROMPT,
        11,
    ],
]


def round_to_nearest_resolution_acceptable_by_vae(height, width):
    height = height - (height % pipe.vae_spatial_compression_ratio)
    width = width - (width % pipe.vae_spatial_compression_ratio)
    return height, width


def get_duration(
        prompt,
        nag_negative_prompt, nag_scale,
        height, width, duration_seconds,
        seed, randomize_seed,
        compare,
):
    duration = int(duration_seconds) * 4 + 5
    if compare:
        duration *= 2
    return duration

@spaces.GPU(duration=get_duration)
def generate_video(
        prompt,
        nag_negative_prompt, nag_scale,
        height=DEFAULT_H_SLIDER_VALUE, width=DEFAULT_W_SLIDER_VALUE, duration_seconds=DEFAULT_DURATION_SECONDS,
        seed=DEFAULT_SEED, randomize_seed=False,
        compare=True,
):
    target_h = max(MOD_VALUE, (int(height) // MOD_VALUE) * MOD_VALUE)
    target_w = max(MOD_VALUE, (int(width) // MOD_VALUE) * MOD_VALUE)

    downscaled_height, downscaled_width = int(target_h * DOWNSCALE_FACTOR), int(target_w * DOWNSCALE_FACTOR)
    downscaled_height, downscaled_width = round_to_nearest_resolution_acceptable_by_vae(downscaled_height, downscaled_width)

    num_frames = np.clip(int(round(int(duration_seconds) * FIXED_FPS) + 1), MIN_FRAMES_MODEL, MAX_FRAMES_MODEL)

    current_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    with torch.inference_mode():
        latents = pipe(
            conditions=None,
            prompt=prompt,
            width=downscaled_width,
            height=downscaled_height,
            num_frames=num_frames,
            num_inference_steps=7,
            decode_timestep=0.05,
            guidance_scale=1.0,
            decode_noise_scale=0.025,
            generator=torch.Generator("cuda").manual_seed(current_seed),
            output_type="latent",

            nag_negative_prompt=nag_negative_prompt,
            nag_scale=nag_scale,
        ).frames

        upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
        upscaled_latents = pipe_upsample(
            latents=latents,
            output_type="latent"
        ).frames

        nag_output_frames_list = pipe(
            prompt=prompt,
            width=upscaled_width,
            height=upscaled_height,
            num_frames=num_frames,
            denoise_strength=0.4,  # Effectively, 4 inference steps out of 10
            num_inference_steps=10,
            latents=upscaled_latents,
            decode_timestep=0.05,
            guidance_scale=1.0,
            decode_noise_scale=0.025,
            image_cond_noise_scale=0.025,
            generator=torch.Generator("cuda").manual_seed(current_seed),
            output_type="pil",

            nag_negative_prompt=nag_negative_prompt,
            nag_scale=nag_scale,
        ).frames[0]

    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
        nag_video_path = tmpfile.name
    export_to_video(nag_output_frames_list, nag_video_path, fps=FIXED_FPS)

    if compare:
        latents = pipe(
            conditions=None,
            prompt=prompt,
            width=downscaled_width,
            height=downscaled_height,
            num_frames=num_frames,
            num_inference_steps=7,
            decode_timestep=0.05,
            guidance_scale=1.0,
            decode_noise_scale=0.025,
            generator=torch.Generator("cuda").manual_seed(current_seed),
            output_type="latent",
        ).frames

        upscaled_height, upscaled_width = downscaled_height * 2, downscaled_width * 2
        upscaled_latents = pipe_upsample(
            latents=latents,
            output_type="latent"
        ).frames

        baseline_output_frames_list = pipe(
            prompt=prompt,
            width=upscaled_width,
            height=upscaled_height,
            num_frames=num_frames,
            denoise_strength=0.4,  # Effectively, 4 inference steps out of 10
            num_inference_steps=10,
            latents=upscaled_latents,
            decode_timestep=0.05,
            guidance_scale=1.0,
            decode_noise_scale=0.025,
            image_cond_noise_scale=0.025,
            generator=torch.Generator("cuda").manual_seed(current_seed),
            output_type="pil",
        ).frames[0]

        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpfile:
            baseline_video_path = tmpfile.name
        export_to_video(baseline_output_frames_list, baseline_video_path, fps=FIXED_FPS)
    else:
        baseline_video_path = None

    return nag_video_path, baseline_video_path, current_seed


def generate_video_with_example(
        prompt,
        nag_negative_prompt,
        nag_scale,
):
    nag_video_path, baseline_video_path, seed = generate_video(
        prompt=prompt,
        nag_negative_prompt=nag_negative_prompt, nag_scale=nag_scale,
        height=DEFAULT_H_SLIDER_VALUE, width=DEFAULT_W_SLIDER_VALUE, duration_seconds=DEFAULT_DURATION_SECONDS,
        seed=DEFAULT_SEED, randomize_seed=False,
        compare=True,
    )
    return nag_video_path, baseline_video_path, \
        DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE, \
        DEFAULT_DURATION_SECONDS, seed, True


with gr.Blocks() as demo:
    gr.Markdown('''# Normalized Attention Guidance (NAG) for [LTX Video 0.9.7 Distilled](https://huggingface.co/Lightricks/LTX-Video-0.9.7-distilled)
    NAG demos: [4-Step Wan2.1](https://huggingface.co/spaces/ChenDY/NAG_wan2-1-fast), [FLUX.1-dev](https://huggingface.co/spaces/ChenDY/NAG_FLUX.1-dev), [FLUX.1-schnell](https://huggingface.co/spaces/ChenDY/NAG_FLUX.1-schnell)

    [Paper](https://arxiv.org/abs/2505.21179), [GitHub](https://github.com/ChenDarYen/Normalized-Attention-Guidance), [ComfyUI](https://github.com/ChenDarYen/ComfyUI-NAG)
    
    Implementation of [Normalized Attention Guidance](https://chendaryen.github.io/NAG.github.io/).
    ''')

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                max_lines=3,
                placeholder="Enter your prompt",
            )
            nag_negative_prompt = gr.Textbox(
                label="Negative Prompt for NAG",
                value=DEFAULT_NAG_NEGATIVE_PROMPT,
                max_lines=3,
            )
            nag_scale = gr.Slider(label="NAG Scale", minimum=1., maximum=20., step=0.25, value=11.)
            compare = gr.Checkbox(
                label="Compare with baseline",
                info="If unchecked, only sample with NAG will be generated.", value=True,
            )

            with gr.Accordion("Advanced Settings", open=False):
                duration_seconds_input = gr.Slider(
                    minimum=1, maximum=10, step=1, value=DEFAULT_DURATION_SECONDS,
                    label="Duration (seconds)",
                )
                seed_input = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=DEFAULT_SEED, interactive=True)
                randomize_seed_checkbox = gr.Checkbox(label="Randomize seed", value=True, interactive=True)
                with gr.Row():
                    height_input = gr.Slider(minimum=SLIDER_MIN_H, maximum=SLIDER_MAX_H, step=MOD_VALUE,
                                             value=DEFAULT_H_SLIDER_VALUE,
                                             label=f"Output Height (multiple of {MOD_VALUE})")
                    width_input = gr.Slider(minimum=SLIDER_MIN_W, maximum=SLIDER_MAX_W, step=MOD_VALUE,
                                            value=DEFAULT_W_SLIDER_VALUE,
                                            label=f"Output Width (multiple of {MOD_VALUE})")

            generate_button = gr.Button("Generate Video", variant="primary")
        with gr.Column():
            nag_video_output = gr.Video(label="Video with NAG", autoplay=True, interactive=False)
            baseline_video_output = gr.Video(label="Baseline Video without NAG", autoplay=True, interactive=False)

    gr.Examples(
        examples=examples,
        fn=generate_video_with_example,
        inputs=[prompt, nag_negative_prompt, nag_scale],
        outputs=[
            nag_video_output, baseline_video_output,
            height_input, width_input, duration_seconds_input,
            seed_input,
            compare,
        ],
        cache_examples="lazy"
    )

    ui_inputs = [
        prompt,
        nag_negative_prompt, nag_scale,
        height_input, width_input, duration_seconds_input,
        seed_input, randomize_seed_checkbox,
        compare,
    ]
    generate_button.click(
        fn=generate_video,
        inputs=ui_inputs,
        outputs=[nag_video_output, baseline_video_output, seed_input],
    )

if __name__ == "__main__":
    demo.queue().launch()