#!/usr/bin/env -S uv run --script
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# https://docs.astral.sh/uv/guides/scripts/#using-a-shebang-to-create-an-executable-file
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "accelerate",
#   "cosmos-guardrail",
#   "diffusers>=0.34.0",
#   "transformers",
# ]
# [tool.uv]
# exclude-newer = "2025-08-15T00:00:00Z"
# override-dependencies = ["peft>=0.15.0"]
# ///


"""Example of Cosmos-Predict2 Text2Image inference using Hugging Face diffusers."""

import argparse
import pathlib
import textwrap

import diffusers
import torch

ROOT = pathlib.Path(__file__).parents[1]
SEPARATOR = "-" * 20


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("output", type=str, help="Output directory")
    parser.add_argument("--prompt", type=str, required=True, help="Path to prompt text file")
    parser.add_argument("--negative_prompt", type=str, help="Path to negative prompt text file")
    parser.add_argument(
        "--model",
        type=str,
        default="nvidia/Cosmos-Predict2-2B-Text2Image",
        help="Model name or path (Cosmos-Predict2: https://huggingface.co/collections/nvidia/cosmos-predict2-68028efc052239369a0f2959)",
    )
    parser.add_argument("--revision", type=str, help="Model revision (branch name, tag name, or commit id)")
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose output")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
    parser.add_argument(
        "--steps",
        type=int,
        default=35,
        help="The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.",
    )
    parser.add_argument(
        "--guidance",
        type=float,
        default=7,
        help="Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`.",
    )
    parser.add_argument("--height", type=int, default=768, help="The height in pixels of the generated image.")
    parser.add_argument("--width", type=int, default=1360, help="The width in pixels of the generated image.")
    args = parser.parse_args()

    prompt = open(args.prompt).read()
    if args.negative_prompt is not None:
        negative_prompt = open(args.negative_prompt).read()
    else:
        negative_prompt = open(f"{ROOT}/prompts/video2world/negative/default.txt").read()

    output_dir = pathlib.Path(args.output)
    output_dir.mkdir(parents=True, exist_ok=True)

    if args.verbose:
        print(SEPARATOR)
        print("Prompt:")
        print(textwrap.indent(prompt.rstrip(), "  "))
        print("Negative Prompt:")
        print(textwrap.indent(negative_prompt.rstrip(), "  "))
        print(SEPARATOR)

    pipe = diffusers.Cosmos2TextToImagePipeline.from_pretrained(
        args.model,
        revision=args.revision,
        use_safetensors=True,
        torch_dtype=torch.bfloat16,
    )
    pipe.to("cuda")

    print("Generating image...")
    output = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        generator=torch.Generator("cuda").manual_seed(args.seed),
        num_inference_steps=args.steps,
        guidance_scale=args.guidance,
        height=args.height,
        width=args.width,
    ).images[0]
    output.save(output_dir / "output.png")
    print(f"Saved image to: {output_dir / 'output.png'}")


if __name__ == "__main__":
    main()