# SPDX-License-Identifier: Apache-2.0
"""Text to image and an edit on one p100a.

    python examples/quickstart.py "A red fox standing in fresh snow at dawn"
"""
import sys

from tt_flux2_klein import KleinTT, close_device, dram_stats, open_device

prompt = sys.argv[1] if len(sys.argv) > 1 else "A red fox standing in fresh snow at dawn, wildlife photography"

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

    image, timing = model.generate(prompt, 1024, 1024, seed=42)
    image.save("output.png")
    print({k: round(v, 3) for k, v in timing.values.items()})

    # editing: pass 1-3 reference images to the same call
    # from PIL import Image
    # edited, _ = model.generate("Replace the cup with a glass of orange juice.", 512, 512,
    #                            seed=3, images=[Image.open("../media/croissant512.png")])
finally:
    close_device(dev)