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
imageserver
main.py

69
imageserver/main.py
View File

@ -14,8 +14,8 @@ from PIL import Image
load_dotenv() 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 = 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.to("cuda")
class Text(BaseModel): class Text(BaseModel):
text: str text: str
@ -25,39 +25,44 @@ class Text(BaseModel):
app = FastAPI() app = FastAPI()
@app.post("/", @app.post("/",
# responses = { responses = {
# 200: { 200: {
# "content": {"image/png": {}} "content": {"image/png": {}}
# } }
# }, },
# response_class=Response response_class=Response
) )
def root(text: str): async def root(text: str, my_file: UploadFile = File(...)):
prompt = text.replace('+', ' ') prompt = text.replace('+', ' ')
print(prompt) print(prompt)
#request_object_content = file.read() request_object_content = await my_file.read()
# img = Image.open(io.BytesIO(request_object_content)) img = Image.open(io.BytesIO(request_object_content))
# height_orig = img.height height_orig = img.height
# width_orig = img.width width_orig = img.width
# aspect_ratio = width_orig / height_orig aspect_ratio = width_orig / height_orig
# width_new = 512 width_new = 512
# height_new = int(width_new / aspect_ratio) height_new = int(width_new / aspect_ratio)
# img = img.resize((width_new, height_new), 0) 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: imgByteArr = io.BytesIO()
image.save(imgByteArr, format="PNG")
# resp = pipe(prompt, init_image=img) imgByteArr = imgByteArr.getvalue()
# print(resp) running = False
# image = resp.images[0] return Response(content=imgByteArr, media_type="image/png")
# 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")