Save images locally. Add manual seed
This commit is contained in:
parent
2f91297ff9
commit
018163fae8
|
@ -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()
|
||||
|
|
Loading…
Reference in New Issue