Update main.py: Refactor code for CPU compatibility

main.py

@@ -5,12 +5,12 @@ 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 transformers import CLIPFeatureExtractor
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 
 
@@ -18,11 +18,11 @@ from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 async def lifespan(app: FastAPI):
     feature_extractor = CLIPFeatureExtractor.from_pretrained(
         "openai/clip-vit-base-patch32"
-    )
+    ).to("cpu")
 
     safety_checker = StableDiffusionSafetyChecker.from_pretrained(
         "CompVis/stable-diffusion-safety-checker"
-    )
+    ).to("cpu")
 
     text2img = AutoPipelineForText2Image.from_pretrained(
         "stabilityai/sd-turbo",
@@ -39,6 +39,7 @@ async def lifespan(app: FastAPI):
     del inpaint
     del img2img
     del text2img
+
     del safety_checker
     del feature_extractor
 
@@ -68,7 +69,9 @@ async def text_to_image(
     num_inference_steps: int = Form(1),
 ):
     results = request.state.text2img(
-        prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=0.0
+        prompt=prompt,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=0.0,
     )
 
     if not results.nsfw_content_detected[0]: