add flag from mask

app.py

@@ -17,6 +17,7 @@ def process(
     input_image: np.ndarray,
     prompt: str,
     negative_prompt: str,
+    generate_from_mask: bool,
     num_inference_steps: int,
     guidance_scale: float,
     conditioning_scale: float,
@@ -29,6 +30,7 @@ def process(
     #     control_image=input_image,
     #     prompt=prompt,
     #     negative_prompt=negative_prompt,
+    #     generate_from_mask=generate_from_mask,
     #     num_inference_steps=num_inference_steps,
     #     guidance_scale=guidance_scale,
     #     conditioning_scale=conditioning_scale,
@@ -69,6 +71,7 @@ with block:
             input_image = gr.Image(type="numpy")
             prompt = gr.Textbox(label="Prompt")
             negative_prompt = gr.Textbox(label="Negative Prompt")
+            generate_from_mask = gr.Checkbox(label="Input image is already a control mask", value=False)
             run_button = gr.Button(value="Run")
             with gr.Accordion("Advanced options", open=False):
                 target_image_size = gr.Slider(
@@ -107,7 +110,7 @@ with block:
             generated_output = gr.Image(label="Generated", type="numpy", elem_id="generated")
             mask_output = gr.Image(label="Mask", type="numpy", elem_id="mask")
 
-    ips = [input_image, prompt, negative_prompt, num_inference_steps, guidance_scale, conditioning_scale, guess_mode, target_image_size, max_image_size, seed]
+    ips = [input_image, prompt, negative_prompt, generate_from_mask, num_inference_steps, guidance_scale, conditioning_scale, guess_mode, target_image_size, max_image_size, seed]
     run_button.click(fn=process, inputs=ips, outputs=[generated_output, mask_output])
 
 

src/pipeline.py

@@ -41,6 +41,7 @@ class FashionPipeline:
         control_image: np.ndarray,
         prompt: str,
         negative_prompt: str,
+        generate_from_mask: bool,
         num_inference_steps: int,
         guidance_scale: float,
         conditioning_scale: float,
@@ -54,8 +55,11 @@ class FashionPipeline:
         control_image = HWC3(control_image)
 
         # extract segmentation mask
-        segm_mask = generate_mask(control_image, self.segmentation_model, device=self.device)
-        control_mask = self.create_control_mask(segm_mask)
+        if generate_from_mask:
+            control_mask = control_image
+        else:
+            segm_mask = generate_mask(control_image, self.segmentation_model, device=self.device)
+            control_mask = self.create_control_mask(segm_mask)
 
         control_mask = self.adaptive_resize(
             image=control_mask,