speechtoimage
/
speechtoimage
1
0
Fork 0
Code Issues Pull Requests 2 Projects Releases Wiki Activity

Save images locally. Add manual seed

This commit is contained in:
Jimmy 2022-11-01 14:00:35 +13:00
parent 2f91297ff9
commit 018163fae8
1 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
imageserver/main.py
View File

@ -9,6 +9,7 @@ from os import getenv
from fastapi import FastAPI, Response, HTTPException from fastapi import FastAPI, Response, HTTPException
from pydantic import BaseModel from pydantic import BaseModel
import io import io
from PIL.PngImagePlugin import PngInfo
load_dotenv() load_dotenv()
@ -30,16 +31,23 @@ app = FastAPI()
}, },
response_class=Response response_class=Response
) )
def root(text: Text): def root(text: str):
prompt = text.text prompt = text.replace('+', ' ')
print(prompt) print(prompt)
try: 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: except RuntimeError as e:
raise HTTPException(status_code=202, detail="Busy") raise HTTPException(status_code=202, detail="Busy")
except: except:
raise HTTPException(status_code=504) 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() imgByteArr = io.BytesIO()
image.save(imgByteArr, format="PNG") image.save(imgByteArr, format="PNG")
imgByteArr = imgByteArr.getvalue() imgByteArr = imgByteArr.getvalue()