Update app.py

app.py

@@ -1,89 +1,81 @@
-import cv2
 import torch
-import gc
-import spaces
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
+import os
 import gradio as gr
 import numpy as np
 from PIL import Image
 from PIL.ImageOps import grayscale
-from diffusers import StableDiffusionInpaintPipeline
-
-
-class NecklaceTryOn:
-    def __init__(self):
-        self.model_id = "stabilityai/stable-diffusion-2-inpainting"
-        self.pipeline = StableDiffusionInpaintPipeline.from_pretrained(
-            self.model_id, torch_dtype=torch.float16
-        )
-        self.pipeline = self.pipeline.to("cuda")
-
-    def clear_func(self):
-        torch.cuda.empty_cache()
-        gc.collect()
-
-    @spaces.GPU
-    def clothing_try_on_n_necklace_try_on(self, result_image, mask):
-        """Main method for clothing and necklace try-on."""
-        jewellery_mask = Image.fromarray(
-            np.bitwise_and(np.array(mask), np.array(result_image))
-        )
-        arr_orig = np.array(grayscale(mask))
-
-        result_image = cv2.inpaint(np.array(result_image), arr_orig, 15, cv2.INPAINT_TELEA)
-        result_image = Image.fromarray(result_image)
-
-        arr = arr_orig.copy()
-        mask_y = np.where(arr == arr[arr != 0][0])[0][0]
-        arr[mask_y:, :] = 255
-
-        new_mask = Image.fromarray(arr)
-        mask = new_mask.copy()
-
-        orig_size = result_image.size
-        result_image = result_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 = self.pipeline(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            image=result_image,
-            mask_image=mask,
-            strength=0.95,
-            guidance_score=9,
-        ).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
-        ]
-        self.clear_func()
-        return results[0]
-
-    def launch_interface(self):
-        with gr.Blocks() as interface:
-            with gr.Row():
-                inputImage = gr.Image(label="Input Image", type="pil", image_mode="RGB", interactive=True)
-                mask_image = gr.Image(label="Mask Image", type="pil", image_mode="RGB", interactive=True)
-                outputOne = gr.Image(label="Output", interactive=False)
-
-            submit = gr.Button("Apply")
-
-            submit.click(fn=self.clothing_try_on_n_necklace_try_on, inputs=[inputImage, mask_image],
-                         outputs=[outputOne])
-
-        interface.launch(debug=True)
+import gc
+import spaces
+
+model_id = "stabilityai/stable-diffusion-2-inpainting"
+pipeline = StableDiffusionInpaintPipeline.from_pretrained(
+    model_id, torch_dtype=torch.float16
+)
+pipeline = pipeline.to("cuda")
+
+def clear_func():
+    """Clear GPU memory cache."""
+    torch.cuda.empty_cache()
+    gc.collect()
+
+def process_mask(mask):
+    """Convert mask to binary format (black and white) for inpainting."""
+    mask = mask.convert("L")  # Convert to grayscale
+    mask = np.array(mask)
+
+    # Convert to binary: 0 (black) -> keep, 255 (white) -> modify
+    mask = np.where(mask > 128, 255, 0).astype(np.uint8)
+
+    return Image.fromarray(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
+
+def launch_interface():
+    """Launch the Gradio interface."""
+    with gr.Blocks() as interface:
+        with gr.Row():
+            inputImage = gr.Image(label="Input Image", type="pil", image_mode="RGB", interactive=True)
+            maskImage = gr.Image(label="Input Mask", type="pil", image_mode="RGB", interactive=True)
+            outputOne = gr.Image(label="Output", interactive=False)
+
+        submit = gr.Button("Apply")
+
+        submit.click(fn=clothing_try_on, inputs=[inputImage, maskImage], outputs=[outputOne])
 
+    interface.launch(debug=True)
 
 if __name__ == "__main__":
-    app = NecklaceTryOn()
-    app.launch_interface()
+    launch_interface()