Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import spaces
+import time
 import cv2
 from diffusers import AutoPipelineForInpainting
 from transformers import pipeline
@@ -30,15 +31,10 @@ def get_most_similar_string(target_string, string_array):
             best_match_ratio = similarity_ratio
     return best_match
 
-# Load models
+# Load YOLO model on CPU
 yoloModel = YOLO('yolov8x-seg.pt')
 yoloModel.to('cpu')
 
-sdxl = AutoPipelineForInpainting.from_pretrained(
-    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
-)
-image_captioner = pipeline("image-to-text", model="Abdou/vit-swin-base-224-gpt2-image-captioning")
-
 def getClasses(model, img1):
     results = model([img1])
     out = []
@@ -80,11 +76,14 @@ def joinClasses(classes):
 
 def getSegments(yoloModel, img1):
     classes, image, results1 = getClasses(yoloModel, img1)
+    im = Image.fromarray(image)
+    im.save('classes.jpg')
     allMask = joinClasses(classes)
     return allMask
 
 @spaces.GPU
-def getDescript(image_captioner, img1):
+def getDescript(img1):
+    image_captioner = pipeline("image-to-text", model="Abdou/vit-swin-base-224-gpt2-image-captioning", device='cuda')
     base64_img = image_to_base64(img1)
     caption = image_captioner(base64_img)[0]['generated_text']
     return caption
@@ -105,19 +104,22 @@ def rmGPT(caption, remove_class, change):
         return ' '.join(arstr)
 
 @spaces.GPU
-def ChangeOBJ(sdxl_model, img1, response, mask1):
+def ChangeOBJ(img1, response, mask1):
+    sdxl = AutoPipelineForInpainting.from_pretrained(
+        "diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16
+    ).to('cuda')
     size = img1.size
-    image = sdxl_model(prompt=response, image=img1, mask_image=mask1).images[0]
+    image = sdxl(prompt=response, image=img1, mask_image=mask1).images[0]
     return image.resize((size[0], size[1]))
 
 def full_pipeline(image, target, change):
     img1 = Image.fromarray(image.astype('uint8'), 'RGB')
     allMask = getSegments(yoloModel, img1)
     target_to_remove = get_most_similar_string(target, list(allMask.keys()))
-    caption = getDescript(image_captioner, img1)
+    caption = getDescript(img1)
     response = rmGPT(caption, target_to_remove, change)
     mask1 = allMask[target_to_remove]
-    remimg = ChangeOBJ(sdxl, img1, response, mask1)
+    remimg = ChangeOBJ(img1, response, mask1)
     return remimg, caption, response
 
 iface = gr.Interface(