Add error handling
This commit is contained in:
parent
39f104288b
commit
f76f62e53d
|
@ -6,22 +6,22 @@ import os
|
|||
from diffusers import StableDiffusionPipeline
|
||||
from dotenv import load_dotenv
|
||||
from os import getenv
|
||||
from fastapi import FastAPI, Response
|
||||
from fastapi import FastAPI, Response, HTTPException
|
||||
from pydantic import BaseModel
|
||||
import io
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=getenv("TOKEN"))
|
||||
pipe.to("cuda")
|
||||
|
||||
class Text(BaseModel):
|
||||
text: str
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
@app.get("/",
|
||||
responses = {
|
||||
200: {
|
||||
|
@ -30,20 +30,18 @@ app = FastAPI()
|
|||
},
|
||||
response_class=Response
|
||||
)
|
||||
async def root(text: Text):
|
||||
# get your token at https://huggingface.co/settings/tokens
|
||||
|
||||
def root(text: Text):
|
||||
prompt = text.text
|
||||
print(prompt)
|
||||
image = pipe(prompt).images[0]
|
||||
try:
|
||||
image = pipe(prompt).images[0]
|
||||
except RuntimeError as e:
|
||||
raise HTTPException(status_code=202, detail="Busy")
|
||||
except:
|
||||
raise HTTPException(status_code=504)
|
||||
|
||||
# print(image)
|
||||
|
||||
# image = Image.new('RGB', (1000, 1000), (100,200,10))
|
||||
imgByteArr = io.BytesIO()
|
||||
# image.save expects a file as a argument, passing a bytes io ins
|
||||
image.save(imgByteArr, format="PNG")
|
||||
# Turn the BytesIO object back into a bytes object
|
||||
imgByteArr = imgByteArr.getvalue()
|
||||
# media_type here sets the media type of the actual response sent to the client.
|
||||
running = False
|
||||
return Response(content=imgByteArr, media_type="image/png")
|
||||
|
|
Loading…
Reference in New Issue