rewrite ui

app.py

@@ -9,20 +9,18 @@ from src.pipeline import FashionPipeline, PipelineOutput
 
 
 config = PipelineConfig()
-fashion_pipeline = FashionPipeline(config, device=torch.device('cpu'))
+fashion_pipeline = FashionPipeline(config, device=torch.device('cuda'))
 
 
 def process(input_image: np.ndarray, prompt: str):
-
     output: PipelineOutput = fashion_pipeline(
         control_image=input_image,
         prompt=prompt,
     )
 
     return [
-        output.control_image,
-        output.segmentation_mask,
         output.generated_image,
+        output.segmentation_mask,
     ]
 
 
@@ -60,9 +58,11 @@ with block:
                 """)
 
         with gr.Column():
-            result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery")
+            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]
-    run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
+    run_button.click(fn=process, inputs=ips, outputs=[generated_output, mask_output])
 
 
 block.launch()