File size: 735 Bytes
4430ce1
 
 
 
 
 
a543549
4430ce1
a543549
4430ce1
 
 
 
 
 
 
 
 
 
 
 
a543549
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
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)