fix

app.py

@@ -1,4 +1,5 @@
 import subprocess
+from PIL import Image
 
 def download_file(url, output_filename):
     command = ['wget', '-O', output_filename, '-q', url]
@@ -94,13 +95,12 @@ def remove_background(image_path, mask):
     
     return result
     
-pipe = FluxInpaintPipeline.from_pretrained(bfl_repo, torch_dtype=torch.bfloat16)
+pipe = FluxInpaintPipeline.from_pretrained(bfl_repo, torch_dtype=torch.bfloat16).to(DEVICE)
 MAX_SEED = np.iinfo(np.int32).max
 TRIGGER = "a photo of TOK"
 
 @spaces.GPU(duration=200)
 def execute(image, prompt):
-  pipe.to(DEVICE)
   if not prompt :
         gr.Info("Please enter a text prompt.")
         return None
@@ -123,10 +123,14 @@ def execute(image, prompt):
   for image in range(len(imgs)):
     current_img = imgs[image]
     cv2.imwrite('base_image.jpg', current_img)
-    mask = maskHead('base_image.jpg')
+    cv2.imwrite("mask.jpg", maskHead('base_image.jpg'))
+    
+    im = Image.open('base_image.jpg')
+    mask = Image.open('mask.jpg')
+
     result = pipe(
             prompt=f"{prompt} {TRIGGER}",
-            image=current_img,
+            image=im,
             mask_image=mask,
             width=1024,
             height=1024,