added safety checker

main.py

@@ -1,32 +1,47 @@
-from fastapi import FastAPI, Request, Form, File, UploadFile
+from fastapi import FastAPI, Request, UploadFile, Form, File
 from fastapi.responses import StreamingResponse
 from contextlib import asynccontextmanager
 from starlette.middleware.cors import CORSMiddleware
 
 from PIL import Image
 from io import BytesIO
+from transformers import CLIPFeatureExtractor
 from diffusers import (
     AutoPipelineForText2Image,
     AutoPipelineForImage2Image,
     AutoPipelineForInpainting,
 )
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    text2img = AutoPipelineForText2Image.from_pretrained("stabilityai/sd-turbo").to(
-        "cpu"
+    feature_extractor = CLIPFeatureExtractor.from_pretrained(
+        "openai/clip-vit-base-patch32"
     )
 
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    )
+
+    text2img = AutoPipelineForText2Image.from_pretrained(
+        "stabilityai/sd-turbo",
+        safety_checker=safety_checker,
+        feature_extractor=feature_extractor,
+    ).to("cpu")
+
     img2img = AutoPipelineForImage2Image.from_pipe(text2img).to("cpu")
 
     inpaint = AutoPipelineForInpainting.from_pipe(img2img).to("cpu")
 
     yield {"text2img": text2img, "img2img": img2img, "inpaint": inpaint}
 
-    del text2img
-    del img2img
     del inpaint
+    del img2img
+    del text2img
+
+    del safety_checker
+    del feature_extractor
 
 
 app = FastAPI(lifespan=lifespan)
@@ -49,7 +64,9 @@ async def root():
 
 @app.post("/text-to-image/")
 async def text_to_image(
-    request: Request, prompt: str = Form(...), num_inference_steps: int = Form(1)
+    request: Request,
+    prompt: str = Form(...),
+    num_inference_steps: int = Form(1),
 ):
     image = request.state.text2img(
         prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=0.0
@@ -67,10 +84,11 @@ async def image_to_image(
     prompt: str = Form(...),
     init_image: UploadFile = File(...),
     num_inference_steps: int = Form(2),
-    strength: float = Form(0.5),
+    strength: float = Form(1.0),
 ):
-    bytes = await init_image.read()
-    init_image = Image.open(BytesIO(bytes))
+    init_bytes = await init_image.read()
+    init_image = Image.open(BytesIO(init_bytes))
+    init_width, init_height = init_image.size
     init_image = init_image.convert("RGB").resize((512, 512))
 
     image = request.state.img2img(
@@ -80,6 +98,7 @@ async def image_to_image(
         strength=strength,
         guidance_scale=0.0,
     ).images[0]
+    image = image.resize((init_width, init_height))
 
     bytes = BytesIO()
     image.save(bytes, "PNG")
@@ -93,14 +112,15 @@ async def inpainting(
     prompt: str = Form(...),
     init_image: UploadFile = File(...),
     mask_image: UploadFile = File(...),
-    num_inference_steps: int = Form(3),
-    strength: float = Form(0.5),
+    num_inference_steps: int = Form(2),
+    strength: float = Form(1.0),
 ):
-    bytes = await init_image.read()
-    init_image = Image.open(BytesIO(bytes))
+    init_bytes = await init_image.read()
+    init_image = Image.open(BytesIO(init_bytes))
+    init_width, init_height = init_image.size
     init_image = init_image.convert("RGB").resize((512, 512))
-    bytes = await mask_image.read()
-    mask_image = Image.open(BytesIO(bytes))
+    mask_bytes = await mask_image.read()
+    mask_image = Image.open(BytesIO(mask_bytes))
     mask_image = mask_image.convert("RGB").resize((512, 512))
 
     image = request.state.inpaint(
@@ -111,6 +131,7 @@ async def inpainting(
         strength=strength,
         guidance_scale=0.0,
     ).images[0]
+    image = image.resize((init_width, init_height))
 
     bytes = BytesIO()
     image.save(bytes, "PNG")