add: base64 conversion

app.py

@@ -1,5 +1,4 @@
 import torch
-from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
 import os
 import gradio as gr
 import numpy as np
@@ -8,6 +7,9 @@ from PIL.ImageOps import grayscale
 import gc
 import spaces
 import cv2
+import base64
+from io import BytesIO
+from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
 
 model_id = "stabilityai/stable-diffusion-2-inpainting"
 pipeline = StableDiffusionInpaintPipeline.from_pretrained(
@@ -15,26 +17,26 @@ pipeline = StableDiffusionInpaintPipeline.from_pretrained(
 )
 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 = mask.convert("L")
     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):
     jewellery_mask = Image.fromarray(
-            np.bitwise_and(np.array(mask), np.array(image))
-        )
+        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)
@@ -80,19 +82,41 @@ def clothing_try_on(image, mask):
     clear_func()
     return results[0]
 
-def launch_interface():
-    """Launch the Gradio interface."""
+
+def base64_to_image(base64_str):
+    image_data = base64.b64decode(base64_str)
+    image = Image.open(BytesIO(image_data))
+    return image
+
+
+def image_to_base64(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
+def clothing_try_on_base64(input_image_base64, mask_image_base64):
+    image = base64_to_image(input_image_base64)
+    mask = base64_to_image(mask_image_base64)
+
+    output_image = clothing_try_on(image, mask)
+
+    return image_to_base64(output_image)
+
+
+def launch_interface_base64():
     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)
+            inputImage = gr.Textbox(label="Input Image (Base64)", lines=4)
+            maskImage = gr.Textbox(label="Input Mask (Base64)", lines=4)
+            outputOne = gr.Textbox(label="Output (Base64)", lines=4)
 
         submit = gr.Button("Apply")
 
-        submit.click(fn=clothing_try_on, inputs=[inputImage, maskImage], outputs=[outputOne])
+        submit.click(fn=clothing_try_on_base64, inputs=[inputImage, maskImage], outputs=[outputOne])
+
+        interface.launch(debug=True)
 
-    interface.launch(debug=True)
 
 if __name__ == "__main__":
-    launch_interface()
+    launch_interface_base64()