Bug Fix in Sketch affecting other areas

README.md

@@ -4,7 +4,7 @@ emoji: 🌖
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-python_version: 3.12.8
+python_version: 3.12.3
 sdk_version: 5.22.0
 app_file: app.py
 pinned: false
@@ -24,7 +24,7 @@ thumbnail: >-
 
 # Hex Game Maker
 ## Description
-Welcome to Hex Game Maker, the ultimate tool for transforming your images into mesmerizing hexagon grid masterpieces! Well, this was a test project for HexaGrid. It has some nice features that did not make it into the final version of the program.
+Welcome to Hex Game Maker, the ultimate tool for transforming your images into mesmerizing hexagon grid masterpieces! **Well, this is a test project for HexaGrid.** It has some nice features that did not make it into the final version of the program.
 The intention was to do a full conversion, but the limitation on Negative Prompts is what killed this approach. It does not consistently render table top maps.. but it can do a lot!
 
 ### What Can You Do?

app.py

@@ -54,6 +54,7 @@ from modules.constants import (
     default_lut_example_img, 
     lut_files, 
     MAX_SEED, 
+    IS_SHARED_SPACE,
     # lut_folder,cards, 
     # cards_alternating, 
     # card_colors, 
@@ -348,9 +349,14 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
     flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()
     if flash_attention_enabled:
         pipe.attn_implementation="flash_attention_2"
-    # Compile UNet
-    #pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")
-    pipe.vae.enable_tiling()  # For larger resolutions if needed
+
+    if IS_SHARED_SPACE:
+        pipe.vae.enable_tiling()  # For larger resolutions if needed
+    else:
+        # Compile UNet
+        #pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")
+        #pipe.enable_model_cpu_offload() #for smaller GPUs
+        pipe.vae.enable_slicing()
 
     # Disable unnecessary features
     pipe.safety_checker = None
@@ -384,12 +390,18 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()
     if flash_attention_enabled:
         pipe_i2i.attn_implementation="flash_attention_2"
-    # Compile UNet
-    #pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead")
-    pipe.vae.enable_tiling()  # For larger resolutions if needed
+
+    if IS_SHARED_SPACE:
+         pipe_i2i.vae.enable_tiling()  # For larger resolutions if needed
+    else:
+        # Compile UNet
+        # pipe.transformer = torch.compile(pipe.transformer, mode="reduce-overhead") # uses the other pipe's transformer
+        # pipe_i2i.enable_model_cpu_offload() #for smaller GPUs
+         pipe_i2i.vae.enable_slicing()
+
 
     # Disable unnecessary features
-    pipe.safety_checker = None
+    pipe_i2i.safety_checker = None
     image_input = open_image(image_input_path)
     print(f"\nGenerating image with prompt: {prompt_mash} and {image_input_path}\n")
     approx_tokens= approximate_token_count(prompt_mash)
@@ -399,22 +411,24 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     else:
         prompt = prompt_mash
         prompt2 = None
-    final_image = pipe_i2i(
-        prompt=prompt,
-        prompt_2=prompt2,
-        image=image_input,
-        strength=image_strength,
-        num_inference_steps=steps,
-        guidance_scale=cfg_scale,
-        width=width,
-        height=height,
-        generator=generator,
-        joint_attention_kwargs={"scale": lora_scale},
-        output_type="pil",
-    ).images[0]
+    with calculateDuration("Generating image"):
+        # Generate image
+        final_image = pipe_i2i(
+            prompt=prompt,
+            prompt_2=prompt2,
+            image=image_input,
+            strength=image_strength,
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator,
+            joint_attention_kwargs={"scale": lora_scale},
+            output_type="pil",
+        ).images[0]
     return final_image
 
-@spaces.GPU(duration=140)
+@spaces.GPU(duration=140,progress=gr.Progress(track_tqdm=True))
 def run_lora(prompt, map_option, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, enlarge, use_conditioned_image=False, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.🧨")
@@ -422,6 +436,8 @@ def run_lora(prompt, map_option, image_input, image_strength, cfg_scale, steps,
     # handle selecting a conditioned image from the gallery
     global current_prerendered_image
     conditioned_image=None
+    formatted_map_option = map_option.lower().replace(' ', '_')
+
     if use_conditioned_image:
         print(f"Conditioned path: {current_prerendered_image.value}.. converting to RGB\n")
         # ensure the conditioned image is an image and not a string, cannot use RGBA
@@ -468,15 +484,18 @@ def run_lora(prompt, map_option, image_input, image_strength, cfg_scale, steps,
             
     if(image_input is not None):
         print(f"\nGenerating image to image with seed: {seed}\n")
-        final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed, progress)
+        generated_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed, progress)
+
         if enlarge:
-            upscaled_image = upscale_image(final_image, max(1.0,min((TARGET_SIZE[0]/width),(TARGET_SIZE[1]/height))))
-            # Save the upscaled image to a temporary file
-            with NamedTemporaryFile(delete=False, suffix=".png") as tmp_upscaled:
-                upscaled_image.save(tmp_upscaled.name, format="PNG")
-                temp_files.append(tmp_upscaled.name)
-                print(f"Upscaled image saved to {tmp_upscaled.name}")
-                final_image = tmp_upscaled.name
+            upscaled_image = upscale_image(generated_image, max(1.0,min((TARGET_SIZE[0]/width),(TARGET_SIZE[1]/height))))
+        else:
+            upscaled_image = generated_image
+        # Save the upscaled image to a temporary file
+        with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_") as tmp_upscaled:
+            upscaled_image.save(tmp_upscaled.name, format="PNG")
+            temp_files.append(tmp_upscaled.name)
+            print(f"Upscaled image saved to {tmp_upscaled.name}")
+            final_image = tmp_upscaled.name
         yield final_image, seed, gr.update(visible=False)
     else:
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
@@ -485,18 +504,20 @@ def run_lora(prompt, map_option, image_input, image_strength, cfg_scale, steps,
         step_counter = 0
         for image in image_generator:
             step_counter+=1
-            final_image = image
+            generated_image = image
             progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
             yield image, seed, gr.update(value=progress_bar, visible=True)
 
         if enlarge:
-            upscaled_image = upscale_image(final_image, max(1.0,min((TARGET_SIZE[0]/width),(TARGET_SIZE[1]/height))))
-            # Save the upscaled image to a temporary file
-            with NamedTemporaryFile(delete=False, suffix=".png") as tmp_upscaled:
-                upscaled_image.save(tmp_upscaled.name, format="PNG")
-                temp_files.append(tmp_upscaled.name)
-                print(f"Upscaled image saved to {tmp_upscaled.name}")
-                final_image = tmp_upscaled.name
+            upscaled_image = upscale_image(generated_image, max(1.0,min((TARGET_SIZE[0]/width),(TARGET_SIZE[1]/height))))
+        else:
+            upscaled_image = generated_image
+        # Save the upscaled image to a temporary file
+        with NamedTemporaryFile(delete=False, suffix=".png", prefix=f"{formatted_map_option}_") as tmp_upscaled:
+            upscaled_image.save(tmp_upscaled.name, format="PNG")
+            temp_files.append(tmp_upscaled.name)
+            print(f"Upscaled image saved to {tmp_upscaled.name}")
+            final_image = tmp_upscaled.name
         yield final_image, seed, gr.update(value=progress_bar, visible=False)
         
 def get_huggingface_safetensors(link):
@@ -631,9 +652,11 @@ def update_sketch_dimensions(input_image, sketch_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:
+    if (in_img) and (in_img.size != sk_img.size):
         sk_img = sk_img.resize(in_img.size, Image.LANCZOS)
-    return sk_img
+        return [sk_img, gr.update(width=in_img.width, height=in_img.height)]
+    else:
+        return [sk_img, gr.update()]
 
 @spaces.GPU()
 def getVersions():
@@ -722,7 +745,7 @@ with gr.Blocks(css_paths="style_20250314.css", title=title, theme='Surn/beeuty',
                     lut_intensity = gr.Slider(label="Filter Intensity", minimum=-200, maximum=200, value=100, info="0 none, negative inverts the filter", interactive=True)
                     apply_lut_button = gr.Button("Apply Filter (LUT)", elem_classes="solid", elem_id="apply_lut_button")
                 with gr.Row():
-                    lut_example_image = gr.Image(type="pil", label="Filter (LUT) Example Image", value=default_lut_example_img)
+                    lut_example_image = gr.Image(type="pil", label="Filter (LUT) Example Image", value=default_lut_example_img, format="png")
                 with gr.Row():
                     with gr.Column():
                         gr.Markdown("""
@@ -933,7 +956,7 @@ with gr.Blocks(css_paths="style_20250314.css", title=title, theme='Surn/beeuty',
     ).then(
         fn=update_sketch_dimensions,
         inputs=[input_image, sketch_image],
-        outputs=[sketch_image]
+        outputs=[sketch_image, sketch_image]
     )
     prerendered_image_gallery.select(
         fn=on_prerendered_gallery_selection, 
@@ -960,7 +983,7 @@ with gr.Blocks(css_paths="style_20250314.css", title=title, theme='Surn/beeuty',
     ).then(
         fn=update_sketch_dimensions,
         inputs=[input_image, sketch_image],
-        outputs=[sketch_image]
+        outputs=[sketch_image, sketch_image]
     )
     lora_gallery.select(
         update_selection,
@@ -984,7 +1007,7 @@ with gr.Blocks(css_paths="style_20250314.css", title=title, theme='Surn/beeuty',
     ).then(
         fn=update_sketch_dimensions,
         inputs=[input_image, sketch_image],
-        outputs=[sketch_image]
+        outputs=[sketch_image, sketch_image]
     )
 
 load_env_vars(dotenv_path)