Textwizai / image_generation.py
Erfan11's picture
Create image_generation.py
0602100 verified
raw
history blame
566 Bytes
from PIL import Image
import torch
from transformers import CLIPProcessor, CLIPModel
def generate_image_from_text(prompt):
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
inputs = processor(text=[prompt], images=None, return_tensors="pt", padding=True)
outputs = model(**inputs)
# Dummy image return (you can integrate real image generation later)
img = Image.new("RGB", (256, 256), color="blue") # Just a placeholder image
return img