# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.

import argparse
import os
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, PreTrainedTokenizerFast

# Keep behavior aligned with the existing demo entrypoints.
os.environ.setdefault("DEBUG_FIX_PADDING", "1")
os.environ.setdefault("NOT_ALWASY_DO_2DPOOL", "1")

DEFAULT_CKPT = "/path/to/nemotron-labs-diffuion-image-8B-release"
DEFAULT_PROMPT = (
    "A full-body shot of hyper-realistic female cyborg, human facial skin seamlessly integrated with a glossy white mechanical head shell. "
    "Features a realistic human ear, blue eyes. bright, outdoor, background with blue sky, illuminated by striking bright white studio lighting, "
    "casting soft shadows. Cyberpunk aesthetic, high-tech minimalism, shot on 85mm lens, photorealistic, Unreal Engine 5 render, intricately detailed, "
    "8k resolution, high dynamic range, chest with whit armor plate, cute, beautiful, sexy, glossy surface, reflective, Artstation, pixiv, no hair, "
    "3D render, stylized eyesz"
)
DEFAULT_OUTPUT = "/path/to/output/demo_inference_release.webp"

SCHEDULE_CHOICES = ["shift"]
CONFIDENCE_POLICY_CHOICES = ["mask_git", "mmada", "stratified"]
SCHEDULE_TEMP_CHOICES = ["linear", "cosine2", "shift", "exp"]
RESOLUTION_CHOICES = [256, 512, 1024]

# Match the default generation setup in gradio_t2i_demo.py.
DEFAULT_GENERATION_CONFIG = {
    "guidance_scale": 5.0,
    "n_steps": 64,
    "shift": 5,
    "schedule": "shift",
    "alg_temp": 1.0,
    "dynamic_temperature": False,
    "min_temperature": 0.01,
    "schedule_temp": "linear",
    "temperature": 0.86,
    "confidence_policy": "mmada",
    "micro_cond": "ORIGINAL WIDTH : 1024; ORIGINAL HEIGHT : 1024; TOP : 0; LEFT : 0; SCORE : 6.520; HPS: 3.220",
    "template": "Generate an image with the caption:\n <prompt>",
    "edit_threshold": 0.6,
}


def load_release_model_and_tokenizer(model_path: str, device: str):
    tokenizer = PreTrainedTokenizerFast.from_pretrained(model_path)
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        low_cpu_mem_usage=False,
    )
    model.to(device)
    model.eval()
    model.requires_grad_(False)

    return tokenizer, model


def n_tokens_from_resolution(image_resolution: int) -> int:
    return (image_resolution // 16) * (image_resolution // 16)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Single-image LaVida-O text-to-image inference using the release package defaults.",
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument(
        "--pretrained",
        type=str,
        default=DEFAULT_CKPT,
        help="Path to the model directory.\nChoices: any local HF-style checkpoint or release directory.",
    )
    parser.add_argument(
        "--prompt",
        type=str,
        default=DEFAULT_PROMPT,
        help="Prompt text for generation.\nChoices: any text string.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default=DEFAULT_OUTPUT,
        help="Output image path.\nChoices: any writable file path; extension should match a Pillow-supported format such as .webp, .png, or .jpg.",
    )
    parser.add_argument(
        "--image-resolution",
        type=int,
        choices=RESOLUTION_CHOICES,
        default=1024,
        help="Output resolution in pixels.\nChoices: 256, 512, 1024.",
    )
    parser.add_argument(
        "--guidance-scale",
        type=float,
        default=DEFAULT_GENERATION_CONFIG["guidance_scale"],
        help="Classifier-free guidance strength.\nChoices: any positive float; default is the inline release demo setting.",
    )
    parser.add_argument(
        "--temperature",
        type=float,
        default=DEFAULT_GENERATION_CONFIG["temperature"],
        help="Sampling temperature for token draws.\nChoices: any positive float; lower is more conservative.",
    )
    parser.add_argument(
        "--n-steps",
        type=int,
        default=DEFAULT_GENERATION_CONFIG["n_steps"],
        help="Number of denoising steps.\nChoices: any positive integer; default is the inline release demo setting.",
    )
    parser.add_argument(
        "--schedule",
        type=str,
        choices=SCHEDULE_CHOICES,
        default=DEFAULT_GENERATION_CONFIG["schedule"],
        help="Token transfer schedule.\nChoices: shift.",
    )
    parser.add_argument(
        "--shift",
        type=int,
        default=DEFAULT_GENERATION_CONFIG["shift"],
        help="Shift parameter used by the shift schedule.\nChoices: any non-negative integer; default is the inline release demo setting.",
    )
    parser.add_argument(
        "--confidence-policy",
        type=str,
        choices=CONFIDENCE_POLICY_CHOICES,
        default=DEFAULT_GENERATION_CONFIG["confidence_policy"],
        help="Policy for selecting which masked tokens to reveal next.\nChoices: mask_git, mmada, stratified.",
    )
    parser.add_argument(
        "--schedule-temp",
        type=str,
        choices=SCHEDULE_TEMP_CHOICES,
        default=DEFAULT_GENERATION_CONFIG["schedule_temp"],
        help="Temperature schedule shape across denoising steps.\nChoices: linear, cosine2, shift, exp.",
    )
    parser.add_argument(
        "--alg-temp",
        type=float,
        default=DEFAULT_GENERATION_CONFIG["alg_temp"],
        help="Confidence-ranking temperature used by the reveal policy.\nChoices: any positive float; default is the inline release demo setting.",
    )
    parser.add_argument(
        "--dynamic-temperature",
        action="store_true",
        default=DEFAULT_GENERATION_CONFIG["dynamic_temperature"],
        help="Enable dynamic temperature scaling over time.\nChoices: enabled with --dynamic-temperature, disabled with --no-dynamic-temperature.",
    )
    parser.add_argument(
        "--no-dynamic-temperature",
        dest="dynamic_temperature",
        action="store_false",
        help="Disable dynamic temperature scaling.",
    )
    parser.add_argument(
        "--edit-threshold",
        type=float,
        default=DEFAULT_GENERATION_CONFIG["edit_threshold"],
        help="Post-sampling token editing threshold.\nChoices: any float in practice; use -1 to disable edit-based replacement as in the Gradio demo semantics.",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed.\nChoices: any integer; use -1 to sample a fresh random seed.",
    )
    parser.add_argument(
        "--micro-cond",
        type=str,
        default=DEFAULT_GENERATION_CONFIG["micro_cond"],
        help="Micro-conditioning string injected into the prompt template.\nChoices: any string matching the model's expected metadata style.",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="Torch device for model execution.\nChoices: typically cuda or cpu; cuda is expected for practical inference speed.",
    )
    parser.add_argument(
        "--use-cache",
        action="store_true",
        help="Enable KV-cache prefill path during denoising.\nChoices: set flag to enable, omit to disable.",
    )
    parser.add_argument(
        "--is-legacy",
        action="store_true",
        help="Use legacy generation behavior expected by older checkpoints.\nChoices: set flag to enable, omit to disable.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    gen_cfg = dict(DEFAULT_GENERATION_CONFIG)
    gen_cfg.update(
        micro_cond=args.micro_cond,
        guidance_scale=args.guidance_scale,
        temperature=args.temperature,
        edit_threshold=args.edit_threshold,
        n_steps=args.n_steps,
        schedule=args.schedule,
        shift=args.shift,
        confidence_policy=args.confidence_policy,
        schedule_temp=args.schedule_temp,
        alg_temp=args.alg_temp,
        dynamic_temperature=args.dynamic_temperature,
        block_policy=2,
    )

    if args.seed < 0:
        args.seed = int(torch.seed() % (2**31 - 1))
    torch.manual_seed(args.seed)

    tokenizer, model = load_release_model_and_tokenizer(args.pretrained, args.device)
    model.config.dlm_paradigm = "bidirectional"

    with torch.no_grad():
        with torch.inference_mode():
            image = model.text_to_image(
                args.prompt,
                tokenizer=tokenizer,
                **gen_cfg,
                image_resolution=args.image_resolution,
                n_tokens=n_tokens_from_resolution(args.image_resolution),
                is_legacy=args.is_legacy,
                use_cache=args.use_cache,
                disable_tqdm=False,
                return_intermediate_steps=False,
            )

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    image.save(output_path)

    print(f"Saved image to {output_path}")
    print(f"Seed: {args.seed}")
    print(f"Resolution: {args.image_resolution}")
    print(f"Checkpoint: {args.pretrained}")


if __name__ == "__main__":
    main()