Update app.py

app.py

@@ -323,12 +323,9 @@ div#component-17 {
 }
 
 
-@media only screen and (max-width: 600px) {
-    .gr-row {
-        flex-direction: column;
-    }
-    .gr-column {
-        margin-bottom: 10px;
+@media screen and (max-width: 600px) {
+    .img-row{
+        display: block !important;
     }
 }
 
@@ -368,11 +365,12 @@ with gr.Blocks(css=css, fill_height=True) as demo:
                             paste_back = gr.Checkbox(True, label="Paste back original")
                 with gr.Row(equal_height=False):
                     input_image = gr.ImageMask(
-                        type="pil", label="Input Image", layers=True
+                        type="pil", label="Input Image", layers=True, elem_classes="img-row"
                     )
                     result = ImageSlider(
                         interactive=False,
-                        label="Generated Image",
+                        label="Generated Image", 
+                        elem_classes="img-row"
                     )
                 use_as_input_button = gr.Button("Use as Input Image", visible=False)
                 use_as_input_button.click(