Update app.py

app.py

@@ -25,6 +25,10 @@ from huggingface_hub import (
     ModelCard,
     snapshot_download)
 
+import torch
+from diffusers import FluxControlNetPipeline
+from diffusers import FluxControlNetModel
+
 import spaces
 
 #---if env = local or colab---
@@ -1430,6 +1434,14 @@ pipe_i2i = AutoPipelineForImage2Image.from_pretrained(base_model,
                                                       torch_dtype=dtype
                                                      )
 
+
+# CONTROLNET
+controlnet_model = "InstantX/FLUX.1-dev-controlnet-canny"
+controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
+pipe_canny = FluxControlNetPipeline.from_pretrained(
+    base_model, controlnet=controlnet, torch_dtype=torch.bfloat16
+)
+
 MAX_SEED = 2**32-1
 
 pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)
@@ -1473,6 +1485,25 @@ def update_selection(evt: gr.SelectData, width, height):
         height,
     )
 
+import cv2
+def generate_canny(image, type="canny"):
+    img = cv2.imread(image)
+    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img_blur = cv2.GaussianBlur(img_gray, (3,3), 0) 
+    
+    if(type == "canny"):
+        edges = cv2.Canny(image=img_blur, threshold1=100, threshold2=200
+        cv2.imwrite('output.jpg', edges)
+        return 'output.jpg'
+
+    if(type == "sobel"):
+        sobelxy = cv2.Sobel(src=img_blur, ddepth=cv2.CV_64F, dx=1, dy=1, ksize=5)
+        cv2.imwrite('output.jpg', sobelxy)
+        return 'output.jpg'
+
+    return image
+    
+
 @spaces.GPU(duration=100)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress):
     pipe.to("cuda")
@@ -1511,7 +1542,7 @@ def generate_image_to_image(prompt_mash, image_input_path, image_strength, steps
     return final_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, progress=gr.Progress(track_tqdm=True)):
+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:
         raise gr.Error("You must select a LoRA before proceeding.🧨")
     selected_lora = loras[selected_index]
@@ -1548,9 +1579,15 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
             seed = random.randint(0, MAX_SEED)
             
     if(image_input is not None):
-        
-        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)
+        if(useCanny):
+            final_image = generate_canny(image_input, "canny")
+            yield final_image, seed, gr.update(visible=False)
+        elif(useSobel):
+            final_image = generate_canny(image_input, "canny")
+            yield 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)
     else:
         image_generator = generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, progress)
     
@@ -1705,6 +1742,8 @@ with gr.Blocks(theme="prithivMLmods/Minecraft-Theme", css=css, delete_cache=(60,
                 with gr.Row():
                     cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28)
+                    useCanny = gr.Checkbox(label="use canny")
+                    useSobel = gr.Checkbox(label="use Sobel")
                 
                 with gr.Row():
                     width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024)
@@ -1732,7 +1771,7 @@ with gr.Blocks(theme="prithivMLmods/Minecraft-Theme", css=css, delete_cache=(60,
     gr.on(
         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],
+        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]
     )