Update app.py

app.py

@@ -35,6 +35,27 @@ bfl_repo="black-forest-labs/FLUX.1-dev"
 BG_COLOR = (255, 255, 255) # white
 MASK_COLOR = (0, 0 , 0) # black
 
+def maskPerson(input):
+    base_options = python.BaseOptions(model_asset_path='selfie_multiclass_256x256.tflite')
+    options = vision.ImageSegmenterOptions(base_options=base_options,
+                                           output_category_mask=True)
+
+    with vision.ImageSegmenter.create_from_options(options) as segmenter:
+        image = mp.Image.create_from_file(input)
+        segmentation_result = segmenter.segment(image)
+        person_mask = segmentation_result.confidence_masks[0]
+
+        image_data = image.numpy_view()
+        fg_image = np.zeros(image_data.shape, dtype=np.uint8)
+        fg_image[:] = MASK_COLOR
+        bg_image = np.zeros(image_data.shape, dtype=np.uint8)
+        bg_image[:] = BG_COLOR
+
+        condition = np.stack((person_mask.numpy_view(),) * 3, axis=-1) > 0.2
+        output_image = np.where(condition, fg_image, bg_image)
+
+        return output_image
+
 def maskHead(input):
   base_options = python.BaseOptions(model_asset_path='selfie_multiclass_256x256.tflite')
   options = vision.ImageSegmenterOptions(base_options=base_options,
@@ -102,8 +123,6 @@ pipe = FluxInpaintPipeline.from_pretrained(bfl_repo, torch_dtype=torch.bfloat16)
 MAX_SEED = np.iinfo(np.int32).max
 TRIGGER = "a photo of TOK"
 
-print(dir(pipe))
-
 
 @spaces.GPU(duration=100)
 def execute(image, prompt, debug=False):
@@ -118,28 +137,49 @@ def execute(image, prompt, debug=False):
   img = cv2.imread(image)
   img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
 
-  imgs = [ random_positioning(img), random_positioning(img)]
+  imgs = [ random_positioning(img)]
 
   pipe.load_lora_weights("XLabs-AI/flux-RealismLora", weight_name='lora.safetensors')
   response = []
 
-  for image in range(len(imgs)):
-    seed_slicer = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed_slicer)
+  seed_slicer = random.randint(0, MAX_SEED)
+  generator = torch.Generator().manual_seed(seed_slicer)
 
+  for image in range(len(imgs)):
     current_img = imgs[image]
     cv2.imwrite('base_image.jpg', current_img)
-    cv2.imwrite("mask.jpg", maskHead('base_image.jpg'))
+    cv2.imwrite("mask_person.jpg", maskPerson('base_image.jpg'))
+    #cv2.imwrite("mask.jpg", maskHead('base_image.jpg'))
     
     im = Image.open('base_image.jpg')
     np_arr = np.array(im)
     rgb_image = cv2.cvtColor(np_arr, cv2.COLOR_BGR2RGB)
+
     im =  Image.fromarray(rgb_image)
-    mask = Image.open('mask.jpg')
+    person = Image.open('mask_person.jpg')
 
     result = pipe(
             prompt=f"{prompt} {TRIGGER}",
             image=im,
+            mask_image=person,
+            width=1024,
+            height=1024,
+            strength=0.85,
+            generator=generator,
+            num_inference_steps=28,
+            max_sequence_length=256,
+            joint_attention_kwargs={"scale": 0.9},
+        ).images[0]
+
+    arr = np.array(result)
+    rgb_image = cv2.cvtColor(arr, cv2.COLOR_BGR2RGB)
+    cv2.imwrite('person.jpg', rgb_image)
+    cv2.imwrite("mask.jpg", maskHead('person.jpg'))
+    mask = Image.open('mask.jpg')
+
+    result = pipe(
+            prompt=f"{prompt} {TRIGGER}",
+            image=result,
             mask_image=mask,
             width=1024,
             height=1024,
@@ -152,8 +192,8 @@ def execute(image, prompt, debug=False):
 
     if debug:
         response.append(im)
+        response.append(person)
         response.append(mask)
-    response.append(result)
 
   return response
 
@@ -173,4 +213,4 @@ iface = gr.Interface(
     outputs="gallery"
 )
 
-iface.launch()
+iface.launch(share=True, debug=True)