Clean up

model.py

@@ -111,31 +111,6 @@ class Model:
                          generator=generator,
                          image=control_image).images
 
-    def process(
-        self,
-        task_name: str,
-        prompt: str,
-        additional_prompt: str,
-        negative_prompt: str,
-        control_image: PIL.Image.Image,
-        vis_control_image: PIL.Image.Image,
-        num_samples: int,
-        num_steps: int,
-        guidance_scale: float,
-        seed: int,
-    ) -> list[PIL.Image.Image]:
-        self.load_controlnet_weight(task_name)
-        results = self.run_pipe(
-            prompt=self.get_prompt(prompt, additional_prompt),
-            negative_prompt=negative_prompt,
-            control_image=control_image,
-            num_images=num_samples,
-            num_steps=num_steps,
-            guidance_scale=guidance_scale,
-            seed=seed,
-        )
-        return [vis_control_image] + results
-
     @staticmethod
     def preprocess_canny(
         input_image: np.ndarray,
@@ -157,7 +132,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         num_steps: int,
         guidance_scale: float,
@@ -171,18 +146,17 @@ class Model:
             low_threshold=low_threshold,
             high_threshold=high_threshold,
         )
-        return self.process(
-            task_name='canny',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('canny')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_hough(
@@ -215,7 +189,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         detect_resolution: int,
         num_steps: int,
@@ -231,18 +205,17 @@ class Model:
             value_threshold=value_threshold,
             distance_threshold=distance_threshold,
         )
-        return self.process(
-            task_name='hough',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('hough')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_hed(
@@ -267,7 +240,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         detect_resolution: int,
         num_steps: int,
@@ -279,18 +252,17 @@ class Model:
             image_resolution=image_resolution,
             detect_resolution=detect_resolution,
         )
-        return self.process(
-            task_name='hed',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('hed')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_scribble(
@@ -311,7 +283,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         num_steps: int,
         guidance_scale: float,
@@ -321,18 +293,17 @@ class Model:
             input_image=input_image,
             image_resolution=image_resolution,
         )
-        return self.process(
-            task_name='scribble',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('scribble')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_scribble_interactive(
@@ -354,7 +325,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         num_steps: int,
         guidance_scale: float,
@@ -364,18 +335,17 @@ class Model:
             input_image=input_image,
             image_resolution=image_resolution,
         )
-        return self.process(
-            task_name='scribble',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('scribble')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_fake_scribble(
@@ -408,7 +378,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         detect_resolution: int,
         num_steps: int,
@@ -420,18 +390,17 @@ class Model:
             image_resolution=image_resolution,
             detect_resolution=detect_resolution,
         )
-        return self.process(
-            task_name='scribble',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('scribble')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_pose(
@@ -462,7 +431,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         detect_resolution: int,
         num_steps: int,
@@ -476,18 +445,17 @@ class Model:
             detect_resolution=detect_resolution,
             is_pose_image=is_pose_image,
         )
-        return self.process(
-            task_name='pose',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('pose')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_seg(
@@ -516,7 +484,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         detect_resolution: int,
         num_steps: int,
@@ -530,18 +498,17 @@ class Model:
             detect_resolution=detect_resolution,
             is_segmentation_map=is_segmentation_map,
         )
-        return self.process(
-            task_name='seg',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('seg')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_depth(
@@ -571,7 +538,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         detect_resolution: int,
         num_steps: int,
@@ -585,18 +552,17 @@ class Model:
             detect_resolution=detect_resolution,
             is_depth_image=is_depth_image,
         )
-        return self.process(
-            task_name='depth',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('depth')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results
 
     @staticmethod
     def preprocess_normal(
@@ -628,7 +594,7 @@ class Model:
         prompt: str,
         additional_prompt: str,
         negative_prompt: str,
-        num_samples: int,
+        num_images: int,
         image_resolution: int,
         detect_resolution: int,
         num_steps: int,
@@ -644,15 +610,14 @@ class Model:
             bg_threshold=bg_threshold,
             is_normal_image=is_normal_image,
         )
-        return self.process(
-            task_name='normal',
-            prompt=prompt,
-            additional_prompt=additional_prompt,
+        self.load_controlnet_weight('normal')
+        results = self.run_pipe(
+            prompt=self.get_prompt(prompt, additional_prompt),
             negative_prompt=negative_prompt,
             control_image=control_image,
-            vis_control_image=vis_control_image,
-            num_samples=num_samples,
+            num_images=num_images,
             num_steps=num_steps,
             guidance_scale=guidance_scale,
             seed=seed,
         )
+        return [vis_control_image] + results