import argparse
from pathlib import Path

import torch
from diffusers import QwenImage21Pipeline


PROMPT = (
    "A moonlit fantasy garden scene at night fills the frame in a highly detailed pixel-art style. "
    "The upper portion shows a deep navy-to-deep-blue sky with a subtle gradient, scattered with many"
    " small square stars in white, pale yellow, and icy blue. Soft, blocky clouds stretch across the "
    "left and right horizon areas, rendered with layered blue and gray pixels that create a gentle "
    "atmospheric depth. In the upper-right quadrant, a very large full moon dominates the sky, "
    "occupying roughly a quarter of the image width. It is circular, warm ivory and pale gold in "
    "color, with visible crater-like mottling, darker gray-brown maria, and a glowing halo of bright "
    "light around its edge. The moonlight casts a cool, silvery illumination over the entire scene.\n\n"
    "Below the sky, the landscape opens into a distant valley. On the right side near the horizon, "
    "dark blue mountain silhouettes rise in layered ridges, with a small castle or town visible at "
    "their base. The castle-like structure has narrow vertical towers and pointed rooftops, glowing "
    "with tiny warm yellow-orange window lights that contrast with the surrounding blue night. "
    "Farther left, darker tree shapes and a fence line appear along the edge of the field, adding "
    "depth and rural context to the setting.\n\nThe central subject is a single adult feminine humanoid"
    " figure with blue skin and long flowing blue hair, shown in full-body view and viewed from a "
    "slightly low angle as if standing among tall flowers. She is positioned slightly left of center,"
    " dancing with one leg lifted and her broad skirt sweeping dramatically to the lower right. Her "
    "posture is graceful and balletic: one arm extends outward toward the lower-left side with "
    "relaxed fingers, while the other arm reaches upward toward the moonlit sky with a delicate open "
    "hand. Her expression is calm and serene, with closed or softly lowered eyes, defined brows, "
    "blue-toned facial features, and a faint composed smile. Her face is stylized rather than "
    "realistic, with smooth blue skin, delicate makeup-like shading, and luminous highlights from the"
    " moonlight.\n\nHer hair is very long, voluminous, and wavy, colored in deep cobalt, royal blue, "
    "and cyan highlights, cascading behind her shoulders and moving as if caught by wind. She wears "
    "ornate silver hair accessories on the right side of her head, including floral or branch-like "
    "metallic ornaments, dangling strands, and a star-shaped earring. Additional silver jewelry "
    "appears at her neck and waist, including delicate chains and pendant-like details. Her clothing "
    "is an elaborate white gown with a fitted bodice, off-shoulder sleeves, ruffled trim, lace-like "
    "edges, and a large layered skirt. The fabric is semi-translucent in places, with soft folds, "
    "cool blue shadows, bright white highlights, and silver embroidery or sparkling accents. The "
    "skirt spreads widely across the lower-right half of the image, forming sweeping arcs of fabric "
    "that overlap the flowers and catch bright moonlit highlights along the outer edges. Her visible "
    "hands and arms are slender, blue-skinned, and decorated with fine silver bracelets and arm "
    "jewelry.\n\nThe foreground is filled with dense white roses and dark green foliage. Large rose "
    "blossoms occupy the lower-left, lower-center, and lower-right areas, with additional smaller "
    "blooms and buds scattered throughout the field. The roses have cream-white petals with pale gray"
    " and blue shadowing, yellow centers, and crisp pixelated petal edges. Their leaves are dark "
    "green, blue-green, and almost black in the shadows, with stems and grass rising between them. "
    "The nearest flowers are larger and more detailed, creating a strong sense of depth, while more "
    "distant blossoms blend into the blue field. The color palette is dominated by cool blues, "
    "silvers, whites, and deep greens, punctuated by warm moonlight and tiny golden lights in the "
    "distant settlement.\n\nThe image is composed like a cinematic fantasy illustration, with the "
    "dancer as the central focal point, the oversized moon acting as a dramatic secondary focal "
    "point, and the rose field framing the lower edge. The pixel-art rendering uses visible square "
    "blocks, crisp outlines, and carefully arranged light and shadow rather than smooth painterly "
    "gradients. The overall mood is dreamlike, romantic, magical, and nocturnal, combining a "
    "celestial garden setting with an ethereal dancing figure under radiant moonlight."
)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="Qwen/Qwen-Image-2.1-Turbo")
    parser.add_argument("--output", type=Path, default=Path("comparison"))
    args = parser.parse_args()
    args.output.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(1)

    pipe = QwenImage21Pipeline.from_pretrained(
        args.model,
        revision="fc13ba1986fcb905854dc8de2f52eee395f49b5b",
        torch_dtype=torch.bfloat16,
    )
    pipe.enable_model_cpu_offload()

    with torch.inference_mode():
        for name, tiled in (("untiled", False), ("tiled", True)):
            pipe.vae.disable_tiling()
            if tiled:
                # default value
                pipe.vae.enable_tiling(
                    tile_sample_min_height=256,
                    tile_sample_min_width=256,
                    tile_sample_stride_height=192,
                    tile_sample_stride_width=192,
                )
            image = pipe(
                prompt=PROMPT,
                width=2048,
                height=2048,
                num_inference_steps=8,
                generator=torch.Generator("cuda").manual_seed(42),
            ).images[0]
            image.save(args.output / f"{name}.png")
            print(f"Saved {args.output / (name + '.png')}", flush=True)


if __name__ == "__main__":
    main()