from typing import Optional
import uvicorn
from fastapi import FastAPI, Form, File, UploadFile
from transformers import AutoTokenizer, AutoModelForImageGeneration
from PIL import Image
import io

app = FastAPI()

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/sd-1.5")
model = AutoModelForImageGeneration.from_pretrained("EleutherAI/sd-1.5")

@app.post("/generate_drawing/")
async def generate_drawing(prompt: str = Form(...)):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs['input_ids'], max_length=256, do_sample=True)
    image = Image.open(io.BytesIO(outputs[0].cpu().numpy()))
    return {"image": image}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)