from transformers import AutoModelForGenerativePreTraining
import numpy as np

# Tải model từ file checkpoint
model = AutoModelForGenerativePreTraining.from_pretrained("path/to/your/checkpoint/file.ckpt")

# Mã hóa đầu vào văn bản
input_text = "Tạo một bức tranh đẹp về một cảnh hoàng hôn."
encoded_input = tokenizer.encode(input_text, return_tensors="pt")

# Tạo ảnh
num_images = 1
generated_images = model.generate(encoded_input, num_images=num_images)

# Chuyển đổi đầu ra thành mảng numpy
generated_images = generated_images.numpy()

# Lưu ảnh
for i, image in enumerate(generated_images):
    image = (image * 255).astype(np.uint8)
    image = image.transpose(1, 2, 0)  # Chuyển đổi hình thức của ảnh
    os.makedirs("images", exist_ok=True)
    np.savez("images/image_{}.npz".format(i), image)
    print(f"Ảnh {i} đã được lưu.")

print("Tất cả ảnh đã được tạo và lưu thành công.")