Update app.py

app.py

@@ -1541,6 +1541,28 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     ).images[0]
     return final_image 
 
+def generate_image_canny(prompt_mash, canny, image_strength, steps, cfg_scale, width, height, lora_scale, seed):
+    pipe_canny.to("cuda")
+    control_image = load_image(canny)
+
+    image = pipe(
+        prompt=prompt_mash,
+        control_image=control_image,
+        controlnet_conditioning_scale=0.6,
+        num_inference_steps=28,
+        guidance_scale=3.5,
+        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 image
+
 @spaces.GPU(duration=100)
 def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, useCanny, useSobel, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
@@ -1581,13 +1603,15 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     if(image_input is not None):
         if(useCanny):
             final_image = generate_canny(image_input, "canny")
-            yield final_image, seed, gr.update(visible=False)
+            img = generate_image_canny(prompt_mash, final_image, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
+            yield img, final_image, seed, gr.update(visible=False)
         elif(useSobel):
             final_image = generate_canny(image_input, "sobel")
-            yield final_image, seed, gr.update(visible=False)
+            img = generate_image_canny(prompt_mash, final_image, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
+            yield img, final_image, seed, gr.update(visible=False)
         else: 
             final_image = generate_image_to_image(prompt_mash, image_input, image_strength, steps, cfg_scale, width, height, lora_scale, seed)
-            yield final_image, seed, gr.update(visible=False)
+            yield final_image, final_image, seed, gr.update(visible=False)
     else:
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
     
@@ -1597,9 +1621,9 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             step_counter+=1
             final_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)
+            yield image, image, seed, gr.update(value=progress_bar, visible=True)
             
-        yield final_image, seed, gr.update(value=progress_bar, visible=False)
+        yield final_image,final_image, seed, gr.update(value=progress_bar, visible=False)
         
 def get_huggingface_safetensors(link):
   split_link = link.split("/")
@@ -1732,6 +1756,7 @@ with gr.Blocks(theme="prithivMLmods/Minecraft-Theme", css=css, delete_cache=(60,
         with gr.Column():
             progress_bar = gr.Markdown(elem_id="progress",visible=False)
             result = gr.Image(label="Generated Image")
+            canny = gr.Image(label="Generated Canny")
 
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):
@@ -1772,7 +1797,7 @@ with gr.Blocks(theme="prithivMLmods/Minecraft-Theme", css=css, delete_cache=(60,
         triggers=[generate_button.click, prompt.submit],
         fn=run_lora,
         inputs=[prompt, input_image, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, useCanny, useSobel],
-        outputs=[result, seed, progress_bar]
+        outputs=[result, canny, seed, progress_bar]
     )
 
 app.queue()