diff --git a/imageserver/main.py b/imageserver/main.py index 31d98bf..a8c19d8 100644 --- a/imageserver/main.py +++ b/imageserver/main.py @@ -9,6 +9,7 @@ from os import getenv from fastapi import FastAPI, Response, HTTPException from pydantic import BaseModel import io +from PIL.PngImagePlugin import PngInfo load_dotenv() @@ -30,16 +31,23 @@ app = FastAPI() }, response_class=Response ) -def root(text: Text): - prompt = text.text +def root(text: str): + prompt = text.replace('+', ' ') print(prompt) try: - image = pipe(prompt).images[0] + generator = torch.Generator("cuda").manual_seed(1024) + resp = pipe(prompt) + print(resp) + image = resp.images[0] except RuntimeError as e: raise HTTPException(status_code=202, detail="Busy") - except: + except: raise HTTPException(status_code=504) + metadata = PngInfo() + metadata.add_text("text", prompt) + image.save(f'/images/{str(uuid.uuid4())}.png', pnginfo=metadata) + imgByteArr = io.BytesIO() image.save(imgByteArr, format="PNG") imgByteArr = imgByteArr.getvalue()