jimmy
/
photobooth
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Working image server

This commit is contained in:
Jimmy 2022-12-13 21:23:05 +13:00
parent d261d8e46f
commit 7dba63c954
1 changed files with 37 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

69
imageserver/main.py
View File

@ -14,8 +14,8 @@ from PIL import Image
load_dotenv()
# pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", num_inference_steps=100, revision="fp16", torch_dtype=torch.float16, use_auth_token=getenv("TOKEN"))
# pipe.to("cuda")
pipe = StableDiffusionImg2ImgPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", num_inference_steps=100, revision="fp16", torch_dtype=torch.float16, use_auth_token=getenv("TOKEN"))
pipe.to("cuda")
class Text(BaseModel):
text: str
@ -25,39 +25,44 @@ class Text(BaseModel):
app = FastAPI()
@app.post("/",
# responses = {
# 200: {
# "content": {"image/png": {}}
# }
# },
# response_class=Response
responses = {
200: {
"content": {"image/png": {}}
}
},
response_class=Response
)
def root(text: str):
async def root(text: str, my_file: UploadFile = File(...)):
prompt = text.replace('+', ' ')
print(prompt)
#request_object_content = file.read()
# img = Image.open(io.BytesIO(request_object_content))
request_object_content = await my_file.read()
img = Image.open(io.BytesIO(request_object_content))
# height_orig = img.height
# width_orig = img.width
# aspect_ratio = width_orig / height_orig
# width_new = 512
# height_new = int(width_new / aspect_ratio)
# img = img.resize((width_new, height_new), 0)
height_orig = img.height
width_orig = img.width
aspect_ratio = width_orig / height_orig
width_new = 512
height_new = int(width_new / aspect_ratio)
img = img.resize((width_new, height_new), 0)
try:
resp = pipe(prompt, image=img)
print(resp)
image = resp.images[0]
except RuntimeError as e:
print(e)
raise HTTPException(status_code=202, detail="Busy")
except Exception as e:
raise HTTPException(status_code=504, detail=str(e))
if resp["nsfw_content_detected"] == [True]:
raise HTTPException(status_code=418, detail="NSFW")
# try:
# resp = pipe(prompt, init_image=img)
# print(resp)
# image = resp.images[0]
# except RuntimeError as e:
# print(e)
# raise HTTPException(status_code=202, detail="Busy")
# except:
# raise HTTPException(status_code=504)
imgByteArr = io.BytesIO()
image.save(imgByteArr, format="PNG")
imgByteArr = imgByteArr.getvalue()
running = False
return Response(content=imgByteArr, media_type="image/png")
# imgByteArr = io.BytesIO()
# image.save(imgByteArr, format="PNG")
# imgByteArr = imgByteArr.getvalue()
# running = False
# return Response(content=imgByteArr, media_type="image/png")