add segment_size

app_base.py

@@ -30,6 +30,7 @@ def create_demo() -> gr.Blocks:
         num_steps: int,
         start_step: int,
         guidance_scale: float,
+        segment_size: int,
         generate_size: int,
         upscale_prompt: str,
         upscale_start_size: int = 256,
@@ -72,7 +73,7 @@ def create_demo() -> gr.Blocks:
             upscale_steps=upscale_steps,
         )
         run_task_time, time_cost_str = get_time_cost(run_task_time, time_cost_str)
-        enhanced_image = enhanced_image.resize((generate_size, generate_size))
+        enhanced_image = enhanced_image.resize((segment_size, segment_size))
 
         return enhanced_image, res_image, time_cost_str
 
@@ -106,6 +107,7 @@ def create_demo() -> gr.Blocks:
                 start_step = gr.Slider(minimum=1, maximum=100, value=30, step=1, label="Start Step")
                 with gr.Accordion("Advanced Options", open=False):
                     guidance_scale = gr.Slider(minimum=0, maximum=20, value=0, step=0.5, label="Guidance Scale")
+                    segment_size = gr.Number(label="Segment Size", value=1024)
                     generate_size = gr.Number(label="Generate Size", value=256)
                     mask_expansion = gr.Number(label="Mask Expansion", value=50, visible=True)
                     mask_dilation = gr.Slider(minimum=0, maximum=10, value=2, step=1, label="Mask Dilation")
@@ -128,12 +130,12 @@ def create_demo() -> gr.Blocks:
         
         g_btn.click(
             fn=segment_image,
-            inputs=[input_image, category, generate_size, mask_expansion, mask_dilation],
+            inputs=[input_image, category, segment_size, mask_expansion, mask_dilation],
             outputs=[origin_area_image, croper],
         ).success(
             fn=image_to_image,
             inputs=[origin_area_image, input_image_prompt, edit_prompt,seed,w1, num_steps, start_step,
-                guidance_scale, generate_size, upscale_prompt, upscale_start_size, upscale_steps,
+                guidance_scale, segment_size, generate_size, upscale_prompt, upscale_start_size, upscale_steps,
                 pre_upscale, pre_upscale_start_size, pre_upscale_steps],
             outputs=[enhanced_image, generated_image, generated_cost],
         ).success(