cast image down to 16 bit from 32

JEMeyer · May 16, 2024 · 0e0bfca · 0e0bfca
1 parent 492c062
commit 0e0bfca
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/stablediffusion_fastapi_multigpu/app.py b/stablediffusion_fastapi_multigpu/app.py
@@ -136,7 +136,9 @@ async def img2img(prompt: str = Form(...), image: UploadFile = File(...)):
                 # Ensure to use the selected GPU for computations
                 pipe = img2img_pipes[gpu_id]
 
-                init_image = Image.open(BytesIO(await image.read()))
+                init_image = Image.open(BytesIO(await image.read())).to(
+                    dtype=torch.float16
+                )
 
                 image = pipe(
                     prompt=prompt,