Update model.py

model.py

@@ -25,7 +25,7 @@ from annotator.util import HWC3, resize_image
 
 CONTROLNET_MODEL_IDS = {
 
-    'depth': 'lllyasviel/sd-controlnet-depth',
+    'hough': 'lllyasviel/sd-controlnet-hough',
 
 }
 
@@ -38,7 +38,7 @@ def download_all_controlnet_weights() -> None:
 class Model:
     def __init__(self,
                  base_model_id: str = 'runwayml/stable-diffusion-v1-5',
-                 task_name: str = 'depth'):
+                 task_name: str = 'hough'):
         self.device = torch.device(
             'cuda:0' if torch.cuda.is_available() else 'cpu')
         self.base_model_id = ''
@@ -123,29 +123,32 @@ class Model:
                          generator=generator,
                          image=control_image).images
 
-    @staticmethod
-    def preprocess_depth(
+     @staticmethod
+    def preprocess_hough(
         input_image: np.ndarray,
         image_resolution: int,
         detect_resolution: int,
-        is_depth_image: bool,
+        value_threshold: float,
+        distance_threshold: float,
     ) -> tuple[PIL.Image.Image, PIL.Image.Image]:
         input_image = HWC3(input_image)
-        if not is_depth_image:
-            control_image, _ = apply_midas(
-                resize_image(input_image, detect_resolution))
-            control_image = HWC3(control_image)
-            image = resize_image(input_image, image_resolution)
-            H, W = image.shape[:2]
-            control_image = cv2.resize(control_image, (W, H),
-                                       interpolation=cv2.INTER_LINEAR)
-        else:
-            control_image = resize_image(input_image, image_resolution)
+        control_image = apply_mlsd(
+            resize_image(input_image, detect_resolution), value_threshold,
+            distance_threshold)
+        control_image = HWC3(control_image)
+        image = resize_image(input_image, image_resolution)
+        H, W = image.shape[:2]
+        control_image = cv2.resize(control_image, (W, H),
+                                   interpolation=cv2.INTER_NEAREST)
+
+        vis_control_image = 255 - cv2.dilate(
+            control_image, np.ones(shape=(3, 3), dtype=np.uint8), iterations=1)
+
         return PIL.Image.fromarray(control_image), PIL.Image.fromarray(
-            control_image)
+            vis_control_image)
 
     @torch.inference_mode()
-    def process_depth(
+    def process_hough(
         self,
         input_image: np.ndarray,
         prompt: str,
@@ -157,20 +160,23 @@ class Model:
         num_steps: int,
         guidance_scale: float,
         seed: int,
-        is_depth_image: bool,
+        value_threshold: float,
+        distance_threshold: float,
     ) -> list[PIL.Image.Image]:
-        control_image, vis_control_image = self.preprocess_depth(
+        control_image, vis_control_image = self.preprocess_hough(
             input_image=input_image,
             image_resolution=image_resolution,
             detect_resolution=detect_resolution,
-            is_depth_image=is_depth_image,
+            value_threshold=value_threshold,
+            distance_threshold=distance_threshold,
         )
-        self.load_controlnet_weight('depth')
+        self.load_controlnet_weight('hough')
         results = self.run_pipe(
             prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            num_images=num_images,
+   
+         num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,