File size: 1,007 Bytes
b2e25b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 | # 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)
|