Text_to_image_test / inference.py
hoh2000's picture
Create inference.py
e828422 verified
raw
history blame contribute delete
543 Bytes
from pathlib import Path
from PIL import Image
from transformers import AutoProcessor, AutoModel
import torch
# Load model and processor
model = AutoModel.from_pretrained("cloudqi/cqi_text_to_image_pt_v0", trust_remote_code=True)
processor = AutoProcessor.from_pretrained("cloudqi/cqi_text_to_image_pt_v0", trust_remote_code=True)
def generate_image(prompt):
inputs = processor(prompt, return_tensors="pt")
with torch.no_grad():
output = model(**inputs).images[0] # assumes .images exists
return Image.fromarray(output)