i

app.py

@@ -143,7 +143,8 @@ def predict(
                 )
 
                 rgb_img = tensor_to_pil(vaedecode_sample[0])
-                return flatten([rgb_img]), seed
+                return (rgb_img[0], None, seed)
+
             else:
                 layereddiffusionapply_sample = ld_fg_apply_layered_diffusion(
                     config="SDXL, Conv Injection", weight=1, model=ckpt[0]
@@ -181,7 +182,7 @@ def predict(
                 mask = tensor_to_pil(mask[0])
                 rgb_img = tensor_to_pil(vaedecode_sample[0])
 
-                return flatten([rgba_img, mask]), seed
+                return (rgba_img[0], mask[0], seed)
                 # return flatten([rgba_img, mask, rgb_img, ld_image])
     except Exception as e:
         raise gr.Error(e)
@@ -200,9 +201,6 @@ def predict_examples(
     seed=-1,
     cfg=10,
 ):
-    print(
-        "RUUNING EXAMPLES", prompt, negative_prompt, input_image, remove_bg, cond_mode
-    )
     return predict(
         prompt,
         negative_prompt,
@@ -225,7 +223,8 @@ css = """
 """
 with gr.Blocks(css=css) as blocks:
     gr.Markdown("""# LayerDiffuse (unofficial)
-                Using ComfyUI building blocks with custom node by [huchenlei](https://github.com/huchenlei/ComfyUI-layerdiffuse)
+    Using ComfyUI building blocks with custom node by [huchenlei](https://github.com/huchenlei/ComfyUI-layerdiffuse)  
+    models: [LayerDiffusion/layerdiffusion-v1](https://huggingface.co/LayerDiffusion/layerdiffusion-v1/tree/main)
                 """)
 
     with gr.Row():
@@ -282,8 +281,10 @@ with gr.Blocks(css=css) as blocks:
                     label="Denoise", value=1.0, minimum=0.0, maximum=1.0, step=0.01
                 )
 
-        with gr.Column(scale=1.8):
-            gallery = gr.Gallery(columns=[2], object_fit="contain", height="unset")
+        with gr.Column():
+            image = gr.Image()
+            with gr.Accordion(label="Mask", open=False):
+                mask = gr.Image()
 
     inputs = [
         prompt,
@@ -298,7 +299,7 @@ with gr.Blocks(css=css) as blocks:
         cfg,
         denoise,
     ]
-    outputs = [gallery, curr_seed]
+    outputs = [image, mask, curr_seed]
 
     gr.Examples(
         fn=predict_examples,