Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import os | |
| from model import VirtualStagingToolV2 | |
| def predict(init_image, style, backyard_style, color_preference, additional_info): | |
| vs_tool = VirtualStagingToolV2(diffusion_version="stabilityai/stable-diffusion-2-inpainting") | |
| if backyard_style: | |
| style = backyard_style | |
| output_images, transparent_mask_image = vs_tool.virtual_stage( | |
| image=init_image, style=style, color_preference=color_preference, | |
| additional_info=additional_info, number_images=3) | |
| return output_images[0], output_images[1], output_images[2], transparent_mask_image | |
| image_blocks = gr.Blocks() | |
| with image_blocks as demo: | |
| gr.Markdown( | |
| """ | |
| # Virtual Home Staging | |
| """) | |
| with gr.Group(): | |
| with gr.Box(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| image = gr.Image(source='upload', elem_id="image_upload", | |
| type="pil", label="Upload", | |
| ).style(height=400) | |
| with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True): | |
| style = gr.Dropdown( | |
| ["Bohemian", "Coastal", "Contemporary", "Farmhouse", "French country", | |
| "Glam", "Industrial", "Japandi", "Mid-century modern", "Minimal", "Modern", | |
| "Rustic", "Scandinavian", "Traditional", "Transitional", ], | |
| label="Design theme", elem_id="input-style", | |
| info="only select if the image is NOT a backyard" | |
| ) | |
| backyard_style = gr.Dropdown( | |
| ["Beautiful garden", "Charming playhouse garden landscaping", | |
| "Cottage garden's colorful planting palette", "Cozy corner with fire pit and seating", | |
| "Garden landscaping with gravel landscaping", "Hip california garden landscaping", | |
| "Lush green lawn", "Mediterranean garden landscaping", | |
| "Moss garden", "Outdoor dining and sitting area", | |
| "Party-ready outdoor space with pool, spa, and fire feature", | |
| "Resort-style landscaping and pool", "Round swimming pool with lawn and pool house"], | |
| label="Backyard design theme", elem_id="input-backyard-style", | |
| info="only select if the image is a backyard" | |
| ) | |
| 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, | |
| ) | |
| gr.Markdown("## Image Examples") | |
| gr.Examples( | |
| examples=[os.path.join(os.path.dirname(__file__), | |
| "examples/exciting-small-kitchen-ideas-1821197-hero-d00f516e2fbb4dcabb076ee9685e877a.jpg"), | |
| os.path.join(os.path.dirname(__file__), | |
| "examples/oct-2019-idh-bathroom-reno-ideas-new-gallery-2.jpg"), | |
| os.path.join(os.path.dirname(__file__), | |
| "examples/tips-for-decorating-a-beautiful-bedroom-1976169-hero-e960fbb8311c4b9b875a1813962d34eb.jpg"), | |
| os.path.join(os.path.dirname(__file__), | |
| "examples/backyard_additions.jpg"), | |
| os.path.join(os.path.dirname(__file__), | |
| "examples/living-room-gallery-shelves-l-shaped-couch-ELeyNpyyqpZ8hosOG3EG1X-b5a39646574544e8a75f2961332cd89a.jpg"), | |
| os.path.join(os.path.dirname(__file__), | |
| "examples/modern-dining-room-ideas-4147451-hero-d6333998f8b34620adfd4d99ac732586.jpg"), | |
| ], | |
| inputs=image | |
| ) | |
| with gr.Column(): | |
| mask_image = gr.Image(label="Mask image", elem_id="mask-img", type="pil").style(height=400) | |
| image_out_1 = gr.Image(label="Output 1", elem_id="output-img-1", type="pil").style(height=400) | |
| image_out_2 = gr.Image(label="Output 2", elem_id="output-img-2", type="pil").style(height=400) | |
| image_out_3 = gr.Image(label="Output 3", elem_id="output-img-3", type="pil").style(height=400) | |
| btn.click(fn=predict, inputs=[image, style, backyard_style, color_preference, additional_info], | |
| outputs=[image_out_1, image_out_2, image_out_3, mask_image]) | |
| image_blocks.queue(concurrency_count=3) | |
| image_blocks.launch() |