add additional info

app.py

@@ -4,7 +4,7 @@ import os
 from model import VirtualStagingToolV2
 
 
-def predict(image, style, backyard_style, color_preference):
+def predict(image, style, backyard_style, color_preference, additional_info):
     init_image = image.convert("RGB").resize((512, 512))
 
     vs_tool = VirtualStagingToolV2(diffusion_version="stabilityai/stable-diffusion-2-inpainting")
@@ -12,7 +12,8 @@ def predict(image, style, backyard_style, color_preference):
         style = backyard_style
 
     output_images, transparent_mask_image = vs_tool.virtual_stage(
-        image=init_image, style=style, color_preference=color_preference, number_images=1)
+        image=init_image, style=style, color_preference=color_preference,
+        additional_info=additional_info, number_images=1)
 
     return output_images[0], transparent_mask_image
 
@@ -53,16 +54,20 @@ with image_blocks as demo:
 
                     color_preference = gr.Textbox(placeholder='Enter color preference',
                                                   label="Color preference", elem_id="input-color")
+
+                    additional_info = gr.Textbox(placeholder='Enter additional information',
+                                                  label="Additional information", elem_id="input-add-info")
+
                     btn = gr.Button("Inpaint!").style(
                         margin=False,
                         rounded=(False, True, True, False),
                         full_width=False,
                     )
                 with gr.Column():
-                    mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=800)
-                    image_out = gr.Image(label="Output", elem_id="output-img", type="pil").style(height=800)
+                    mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=500)
+                    image_out = gr.Image(label="Output", elem_id="output-img", type="pil").style(height=500)
 
-            btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference], outputs=[image_out, mask_image])
+            btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference, additional_info], outputs=[image_out, mask_image])
 
             gr.Markdown("## Image Examples")
             gr.Examples(

model.py

@@ -96,7 +96,8 @@ class VirtualStagingToolV2():
             room = 'living room'
         elif 15 in label_ids and 19 in label_ids:
             room = 'dining room'
-
+        else:
+            room ='room'
         label_ids_without_mask = [i for i in label_ids if i not in mask_items]
 
         items = [self.segmentation_model.config.id2label[i] for i in label_ids_without_mask]