#!/usr/bin/env python3
"""Generate a demo image with PixelDiT-T2I-1024."""

from __future__ import annotations

import importlib.util
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent
MODEL_DIR = REPO_ROOT / "PixelDiT-T2I-1024"
OUTPUT_PATH = MODEL_DIR / "demo.png"

DIFFUSERS_SRC = None
for parent in Path(__file__).resolve().parents:
    candidate = parent / "libs/diffusers/src"
    if (candidate / "diffusers/schedulers/flow_dpm.py").is_file():
        DIFFUSERS_SRC = candidate
        break
if DIFFUSERS_SRC is None:
    fallback = Path("/data/projects/Visual-Generative-Foundation-Model-Collection/libs/diffusers/src")
    if (fallback / "diffusers/schedulers/flow_dpm.py").is_file():
        DIFFUSERS_SRC = fallback
if DIFFUSERS_SRC is None:
    raise ImportError("Could not locate libs/diffusers/src for PixelDiT flow DPM sampling.")
if str(DIFFUSERS_SRC) not in sys.path:
    sys.path.insert(0, str(DIFFUSERS_SRC))
import os

os.environ.setdefault("PIXELDIT_DIFFUSERS_SRC", str(DIFFUSERS_SRC))

import torch


def _load_pipeline_class():
    spec = importlib.util.spec_from_file_location("pixeldit_t2i_pipeline", MODEL_DIR / "pipeline.py")
    if spec is None or spec.loader is None:
        raise ImportError(f"Unable to load pipeline from {MODEL_DIR / 'pipeline.py'}")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.PixelDiTT2IPipeline


def main() -> None:
    pipe_cls = _load_pipeline_class()
    pipe = pipe_cls.from_pretrained(
        str(MODEL_DIR),
        local_files_only=True,
        torch_dtype=torch.bfloat16,
    )
    pipe.to("cuda")

    prompt = "A golden retriever playing in a sunny garden"
    print("text_encoder:", type(pipe.text_encoder).__name__)
    print("prompt:", prompt)

    generator = torch.Generator(device="cuda").manual_seed(42)
    image = pipe(
        prompt=prompt,
        negative_prompt="low quality, worst quality, over-saturated, blurry, deformed, watermark",
        height=1024,
        width=1024,
        num_inference_steps=50,
        guidance_scale=2.75,
        use_chi_prompt=True,
        generator=generator,
    ).images[0]
    image.save(OUTPUT_PATH)
    print(f"Saved demo image to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()