from transformers import AutoTokenizer, AutoModelForImageGeneration

tokenizer = AutoTokenizer.from_pretrained("image-gpt2")
model = AutoModelForImageGeneration.from_pretrained("image-gpt2")

text = "a cat sitting on a couch"
input_ids = tokenizer.encode(text, return_tensors="pt")
output = model.generate(input_ids)

# Save the generated image
with open("cat_on_couch.png", "wb") as f:
    f.write(output[0].cpu().numpy())