Add Sketch to system

app.py

@@ -102,8 +102,6 @@ from utils.version_info import (
 )
 #from utils.depth_estimation import (get_depth_map_from_state)
 
-
-
 input_image_palette = []
 current_prerendered_image = gr.State("./images/images/Beeuty-1.png")
 user_dir = constants.TMPDIR
@@ -619,7 +617,43 @@ def add_border(image, mask_width, mask_height, blank_color):
     print(f"Adding border to image with width: {mask_width}, height: {mask_height}, color: {margin_color}")
     return shrink_and_paste_on_blank(bordered_image_output, mask_width, mask_height, margin_color)
 
-
+def on_input_image_change(image_path):
+    if image_path is None:
+        gr.Warning("Please upload an Input Image to get started.")
+        return None, gr.update()
+    img, img_path = convert_to_rgba_png(image_path)
+    with Image.open(img_path) as pil_img:
+        width, height = pil_img.size
+    return [img_path, gr.update(width=width, height=height)]
+
+def update_sketch_dimensions(input_image, sketch_image):
+    # Load the images using open_image() if they are provided as file paths.
+    in_img = open_image(input_image) if isinstance(input_image, str) else input_image
+    sk_img_path, _ = get_image_from_dict(sketch_image)
+    sk_img = open_image(sk_img_path)
+    # Resize sketch image if dimensions don't match input image.
+    if in_img.size != sk_img.size:
+        sk_img = sk_img.resize(in_img.size, Image.LANCZOS)
+    return [sk_img, gr.update(width=in_img.width, height=in_img.height)]
+
+def composite_with_control_sync(input_image, sketch_image, slider_value):
+    # Load the images using open_image() if they are provided as file paths.
+    in_img = open_image(input_image) if isinstance(input_image, str) else input_image
+    sk_img_path, _ = get_image_from_dict(sketch_image)
+    sk_img = open_image(sk_img_path)
+
+    # Resize sketch image if dimensions don't match input image.
+    if in_img.size != sk_img.size:
+        sk_img = sk_img.resize(in_img.size, Image.LANCZOS)
+
+    # Now composite using the original alpha_composite_with_control function.
+    result_img = alpha_composite_with_control(in_img, sk_img, slider_value)
+    return result_img
+
+def replace_input_with_sketch_image(sketch_image):
+    print(f"Sketch Image: {sketch_image}\n")
+    sketch, is_dict = get_image_from_dict(sketch_image)
+    return sketch
 ####################################### DEPTH ESTIMATION #######################################
 
 
@@ -634,7 +668,6 @@ def preprocess_image(image: Image.Image) -> Image.Image:
     processed_image = TRELLIS_PIPELINE.preprocess_image(image)
     return processed_image
 
-
 def pack_state(gs: Gaussian, mesh: MeshExtractResult, name: str) -> dict:
     return {
         'gaussian': {
@@ -651,8 +684,7 @@ def pack_state(gs: Gaussian, mesh: MeshExtractResult, name: str) -> dict:
         },
         'name': name
     }
-    
-    
+
 def unpack_state(state: dict) -> Tuple[Gaussian, edict, str]:
     gs = Gaussian(
         aabb=state['gaussian']['aabb'],
@@ -954,22 +986,35 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
                 elem_classes="centered solid imgcontainer",
                 key="imgInput",
                 image_mode=None,
-                format="PNG"
+                format="PNG",
+                height=512
             )
+            with gr.Accordion("Sketch Pad", open = False, elem_id="sketchpd"):
+                with gr.Row():
+                    sketch_image = gr.Sketchpad(
+                        label="Sketch Image",
+                        type="filepath",
+                        #invert_colors=True,
+                        #sources=['upload','canvas'],
+                        #tool=['editor','select','color-sketch'],
+                        placeholder="Draw a sketch or upload an image. Currently broken in gradio 5.17.1",
+                        interactive=True,
+                        elem_classes="centered solid imgcontainer",
+                        key="imgSketch",
+                        image_mode="RGBA",
+                        format="PNG",
+                        brush=gr.Brush()
+                    )
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        sketch_replace_input_image_button = gr.Button(
+                            "Replace Input Image with sketch",
+                            elem_id="sketch_replace_input_image_button",
+                            elem_classes="solid"
+                        )
+                        sketch_alpha_composite_slider = gr.Slider(0,100,50,0.5, label="Sketch Transparancy", elem_id="alpha_composite_slider")
+                        btn_sketch_alpha_composite = gr.Button("Overlay Sketch on Input Image", elem_id="btn_sketchninput", elem_classes="solid")
 
-        # New code to convert input image to RGBA PNG
-        def on_input_image_change(image_path):
-            if image_path is None:
-                gr.Warning("Please upload an Input Image to get started.")
-                return None
-            img, img_path = convert_to_rgba_png(image_path)
-            return img_path
-
-        input_image.change(
-            fn=on_input_image_change,
-            inputs=[input_image],
-            outputs=[input_image], scroll_to_output=True,
-        )
         with gr.Column():
             with gr.Accordion("Hex Coloring and Exclusion", open = False):
                 with gr.Row():
@@ -1032,7 +1077,7 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
                     outputs=[input_image], 
                     scroll_to_output=True
                 )
-
+            
     with gr.Row():
         with gr.Accordion("Generate AI Image (click here for options)", open = False):
             with gr.Row():
@@ -1267,8 +1312,29 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
         fn=generate_input_image_click,
         inputs=[input_image,map_options, prompt_textbox, negative_prompt_textbox, model_textbox, randomize_seed, seed_slider, gr.State(False), gr.State(0.5), image_size_ratio],
         outputs=[input_image, seed_slider], scroll_to_output=True
+    ).then(
+        fn=update_sketch_dimensions,
+        inputs=[input_image, sketch_image],
+        outputs=[sketch_image, sketch_image]
     )
-
+    input_image.input(
+        fn=on_input_image_change,
+        inputs=[input_image],
+        outputs=[input_image,sketch_image], scroll_to_output=True,
+    )
+    ###################### sketchpad ############################
+    btn_sketch_alpha_composite.click(
+        fn=composite_with_control_sync,
+        inputs=[input_image, sketch_image, sketch_alpha_composite_slider],
+        outputs=[input_image],
+        scroll_to_output=True
+    )
+    sketch_replace_input_image_button.click(
+        lambda sketch_image: replace_input_with_sketch_image(sketch_image),
+        inputs=[sketch_image],
+        outputs=[input_image], scroll_to_output=True
+    )
+    ##################### model #######################################
     model_textbox.change(
         fn=update_prompt_notes,
         inputs=model_textbox,
@@ -1295,6 +1361,10 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
         fn=generate_input_image_click,
         inputs=[input_image, map_options, prompt_textbox, negative_prompt_textbox, model_textbox,randomize_seed, seed_slider, gr.State(True), image_guidance_stength, image_size_ratio],
         outputs=[input_image, seed_slider], scroll_to_output=True
+    ).then(
+        fn=update_sketch_dimensions,
+        inputs=[input_image, sketch_image],
+        outputs=[sketch_image, sketch_image]
     )
 
     # Update the state variable with the prerendered image filepath when an image is selected
@@ -1309,6 +1379,10 @@ with gr.Blocks(css_paths="style_20250128.css", title=title, theme='Surn/beeuty',
         lambda: current_prerendered_image.value,
         inputs=None,
         outputs=[input_image], scroll_to_output=True
+    ).then(
+        fn=update_sketch_dimensions,
+        inputs=[input_image, sketch_image],
+        outputs=[sketch_image, sketch_image]
     )
     output_overlay_composite.change(
         fn=combine_images_with_lerp,

style_20250128.css

@@ -125,4 +125,19 @@ a {
     .gradio-container, .gradio-container::before {
         max-width: 1920px !important;
     }
+}
+
+.sidebar .toggle-button::before {
+    content: 'Sketch Pad';
+    font-weight: bold;
+    transform: rotate(180deg);
+    margin-right: -120px;
+    width: 120px;
+    background-color: rgba(242, 218, 163, 0.62);
+}
+.dark .sidebar .toggle-button::before {
+    background-color: rgba(41, 18, 5, 0.38) !important;
+}
+    .sidebar.open .toggle-button::before {
+    content: '';
 }

utils/file_utils.py

@@ -1,7 +1,7 @@
 # file_utils
 import os
 import utils.constants as constants
-import shutil
+#import shutil
 from pathlib import Path
 
 def cleanup_temp_files():
@@ -11,32 +11,33 @@ def cleanup_temp_files():
         except Exception as e:
             print(f"Failed to delete temp file {file_path}: {e}")
 
-def rename_file_to_lowercase_extension(image_path: str) -> str:
+def rename_file_to_lowercase_extension(file_path: str) -> str:
     """
-    Renames only the file extension to lowercase by copying it to the temporary folder.
+    Renames a file's extension to lowercase in place.
 
     Parameters:
-        image_path (str): The original file path.
+        file_path (str): The original file path.
 
     Returns:
-        str: The new file path in the temporary folder with the lowercase extension.
+        str: The new file path with the lowercase extension.
 
     Raises:
-        Exception: If there is an error copying the file.
+        OSError: If there is an error renaming the file (e.g., file not found, permissions issue).
     """
-    path = Path(image_path)
-    new_suffix = path.suffix.lower()
-    new_path = path.with_suffix(new_suffix)
-
-    # Make a copy in the temporary folder with the lowercase extension
-    if path.suffix != new_suffix:
-        try:
-            shutil.copy(path, new_path)
-            constants.temp_files.append(str(new_path))
-            return str(new_path)
-        except FileExistsError:
-            raise Exception(f"Cannot copy {path} to {new_path}: target file already exists.")
-        except Exception as e:
-            raise Exception(f"Error copying file: {e}")
+    # Split the path into directory and filename
+    directory, filename = os.path.split(file_path)
+    
+    # Split the filename into name and extension
+    name, ext = os.path.splitext(filename)
+    
+    # Convert the extension to lowercase
+    new_ext = ext.lower()
+    
+    # If the extension changes, rename the file
+    if ext != new_ext:
+        new_filename = name + new_ext
+        new_file_path = os.path.join(directory, new_filename)
+        os.rename(file_path, new_file_path)
+        return new_file_path
     else:
-        return str(path)
+        return file_path

utils/image_utils.py

@@ -19,10 +19,10 @@ from utils.file_utils import rename_file_to_lowercase_extension
 
 def get_image_from_dict(image_path):
     if isinstance(image_path, dict) :
-        if 'image' in image_path:
-            image_path = image_path.get('image')
-        elif 'composite' in image_path:
+        if 'composite' in image_path:
             image_path = image_path.get('composite')
+        elif 'image' in image_path:
+            image_path = image_path.get('image')
         else:
             print("\n Unknown image dictionary.\n")
             raise UserWarning("Unknown image dictionary.")