# SPDX-License-Identifier: Apache-2.0
"""Minimal text-to-image run on one p100a.

    python examples/quickstart.py "Three red apples and two green pears on a wooden table"

HiDream-O1 is the undistilled model: 50 steps at CFG 5 means 100 forwards, about 19 s at 1024x1024.
"""
import sys

from tt_hidream_o1 import HiDreamTT, close_device, dram_stats, open_device

prompt = sys.argv[1] if len(sys.argv) > 1 else "Three red apples and two green pears on a wooden table, still life photograph"

dev = open_device()
try:
    model = HiDreamTT(dev)   # downloads the snapshot into the HF cache on first use
    print(f"device DRAM {dram_stats(dev)['allocated_bytes'] / 2**30:.1f} GiB")

    # snap=False generates the size asked for; the official pipeline would round it to ~4 MP.
    # generate() returns the image, the request it resolved, and the timings.
    image, _req, timing = model.generate(prompt, 1024, 1024, seed=1234, snap=False)
    image.save("output.png")
    print({k: round(v, 2) for k, v in timing.values.items()})
finally:
    close_device(dev)