Update app.py

app.py

@@ -31,37 +31,53 @@ def process_mask(mask):
 
 @spaces.GPU
 def clothing_try_on(image, mask):
-    """Perform clothing try-on using the provided image and binary mask."""
-    orig_size = image.size
-
-    # Process and ensure mask is binary
-    mask = process_mask(mask)
-
-    # Resize image and mask for Stable Diffusion
-    image = image.resize((512, 512))
-    mask = mask.resize((512, 512))
-
-    # Prompt and negative prompt
-    prompt = f"South Indian Saree, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple"
-    negative_prompt = "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly"
-
-    # Perform the inpainting using the Stable Diffusion pipeline
-    output = pipeline(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        image=image,
-        mask_image=mask,
-        strength=0.95,
-        guidance_score=9,
-    ).images[0]
-
-    # Resize the output back to the original size
-    output = output.resize(orig_size)
-
-    # Clean GPU memory
-    clear_func()
-
-    return output
+    jewellery_mask = Image.fromarray(
+            np.bitwise_and(np.array(mask), np.array(image))
+        )
+        arr_orig = np.array(grayscale(mask))
+
+        image = cv2.inpaint(np.array(image), arr_orig, 15, cv2.INPAINT_TELEA)
+        image = Image.fromarray(image)
+
+        arr = arr_orig.copy()
+        mask_y = np.where(arr == arr[arr != 0][0])[0][0]
+        arr[mask_y:, :] = 255
+
+        new = Image.fromarray(arr)
+
+        mask = new.copy()
+
+        orig_size = image.size
+
+        image = image.resize((512, 512))
+        mask = mask.resize((512, 512))
+
+        results = []
+        prompt = f" South Indian Saree, properly worn, natural setting, elegant, natural look, neckline without jewellery, simple"
+        negative_prompt = "necklaces, jewellery, jewelry, necklace, neckpiece, garland, chain, neck wear, jewelled neck, jeweled neck, necklace on neck, jewellery on neck, accessories, watermark, text, changed background, wider body, narrower body, bad proportions, extra limbs, mutated hands, changed sizes, altered proportions, unnatural body proportions, blury, ugly"
+
+        output = pipeline(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            mask_image=mask,
+            strength=0.95,
+            guidance_score=9,
+            # generator = torch.Generator("cuda").manual_seed(42)
+        ).images[0]
+
+        output = output.resize(orig_size)
+        temp_generated = np.bitwise_and(
+            np.array(output),
+            np.bitwise_not(np.array(Image.fromarray(arr_orig).convert("RGB"))),
+        )
+        results.append(temp_generated)
+
+        results = [
+            Image.fromarray(np.bitwise_or(x, np.array(jewellery_mask))) for x in results
+        ]
+        clear_func()
+        return results[0]
 
 def launch_interface():
     """Launch the Gradio interface."""