image editor

app.py

@@ -1,15 +1,15 @@
-import sentencepiece
+from typing import Tuple
+from PIL import Image
 import torch
 import spaces
 import gradio as gr
 import os
-from diffusers.pipelines.flux.pipeline_flux_controlnet_inpaint import (
-    FluxControlNetInpaintPipeline,
-)
+# from diffusers.pipelines.flux.pipeline_flux_controlnet_inpaint import (
+#     FluxControlNetInpaintPipeline,
+# )
 from diffusers.pipelines.flux.pipeline_flux_inpaint import FluxInpaintPipeline
-from diffusers.models.controlnet_flux import FluxControlNetModel
+# from diffusers.models.controlnet_flux import FluxControlNetModel
 from controlnet_aux import CannyDetector
-import psutil
 
 # login hf token
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -18,6 +18,7 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 #
 # login()
 
+IMAGE_SIZE = 1024
 dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -34,37 +35,65 @@ pipe.enable_model_cpu_offload()
 
 canny = CannyDetector()
 
+def resize_image_dimensions(
+    original_resolution_wh: Tuple[int, int],
+    maximum_dimension: int = IMAGE_SIZE
+) -> Tuple[int, int]:
+    width, height = original_resolution_wh
 
-def get_system_memory():
-    memory = psutil.virtual_memory()
-    memory_percent = memory.percent
-    memory_used = memory.used / (1024.0**3)
-    memory_total = memory.total / (1024.0**3)
-    return {
-        "percent": f"{memory_percent}%",
-        "used": f"{memory_used:.3f}GB",
-        "total": f"{memory_total:.3f}GB",
-    }
+    # if width <= maximum_dimension and height <= maximum_dimension:
+    #     width = width - (width % 32)
+    #     height = height - (height % 32)
+    #     return width, height
 
+    if width > height:
+        scaling_factor = maximum_dimension / width
+    else:
+        scaling_factor = maximum_dimension / height
 
-@spaces.GPU(duration=60)
+    new_width = int(width * scaling_factor)
+    new_height = int(height * scaling_factor)
+
+    new_width = new_width - (new_width % 32)
+    new_height = new_height - (new_height % 32)
+
+    return new_width, new_height
+
+# def get_system_memory():
+#     memory = psutil.virtual_memory()
+#     memory_percent = memory.percent
+#     memory_used = memory.used / (1024.0**3)
+#     memory_total = memory.total / (1024.0**3)
+#     return {
+#         "percent": f"{memory_percent}%",
+#         "used": f"{memory_used:.3f}GB",
+#         "total": f"{memory_total:.3f}GB",
+#     }
+#
+
+@spaces.GPU(duration=100)
 def inpaint(
     image,
-    mask,
+    # mask,
     prompt,
     strength,
     num_inference_steps,
     guidance_scale,
     controlnet_conditioning_scale,
 ):
-    # canny_image = canny(image)
+
+    image = image['background']
+    mask = image['layers'][0]
+    width, height = resize_image_dimensions(original_resolution_wh=image.size)
+    resized_image = image.resize((width, height), Image.LANCZOS)
+    resized_mask = mask.resize((width, height), Image.LANCZOS)
 
     image_res = pipe(
         prompt,
-        image=image,
+        image=resized_image,
         # control_image=canny_image,
         # controlnet_conditioning_scale=controlnet_conditioning_scale,
-        mask_image=mask,
+        mask_image=resized_mask,
         strength=strength,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
@@ -75,16 +104,23 @@ def inpaint(
 
 with gr.Blocks() as demo:
     # gr.LoginButton()
-    with gr.Row():
-        with gr.Column():
-            gr.Textbox(value="Hello Memory")
-        with gr.Column():
-            gr.JSON(get_system_memory, every=1)
+    # with gr.Row():
+    #     with gr.Column():
+    #         gr.Textbox(value="Hello Memory")
+    #     with gr.Column():
+    #         gr.JSON(get_system_memory, every=1)
     gr.Interface(
         fn=inpaint,
         inputs=[
+            gr.ImageEditor(
+                label='Image',
+                type='pil',
+                sources=["upload", "webcam"],
+                image_mode='RGB',
+                layers=False,
+                brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed")),
             gr.Image(type="pil", label="Input Image"),
-            gr.Image(type="pil", label="Mask Image"),
+            # gr.Image(type="pil", label="Mask Image"),
             gr.Textbox(label="Prompt"),
             gr.Slider(0, 1, value=0.95, label="Strength"),
             gr.Slider(1, 100, value=50, step=1, label="Number of Inference Steps"),

requirements.txt

@@ -9,3 +9,4 @@ tokenizers
 spaces
 huggingface_hub
 psutil
+mediapipe