| import gradio as gr | |
| import requests | |
| from PIL import Image | |
| from io import BytesIO | |
| from diffusers import StableDiffusionImg2ImgPipeline | |
| import torch | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to(device) | |
| def getImg(url): | |
| response = requests.get(url) | |
| img = Image.open(BytesIO(response.content)).convert("RGB") | |
| img.thumbnail((768, 768)) | |
| return img | |
| def staging(init_img, type, style): | |
| prompt = f"interior design ((({style} style))) ((({type}))) with full furnitures, Natural Lighting, Capricious Lighting, Ray Tracing Reflections --ar 16:9 --v 5.2" | |
| negative_prompt = "(((Ugly))), low-resolution, morbid, blurry, cropped, deformed, dehydrated, text, disfigured, duplicate, error, extra arms, extra fingers, extra legs, extra limbs, fused fingers, gross proportions, jpeg artifacts, long neck, low resolution, tiling, poorly drawn feet, extra limbs, disfigured, body out of frame, cut off, low contrast, underexposed, overexposed, bad art, beginner, amateur, distorted face, low quality, lowres, low saturation, deformed body features, watermark, water mark" | |
| result = pipe(prompt, init_img, strength=0.65, guidance_scale=10, negative_prompt=negative_prompt).images[0] | |
| return result | |
| def process(url_or_image, type, style, mode): | |
| if mode == "URL": | |
| init_img = getImg(url_or_image) | |
| else: | |
| init_img = Image.open(url_or_image).convert("RGB") | |
| result_img = staging(init_img, type, style) | |
| return result_img | |
| def update_visibility(input_mode): | |
| if input_mode == "URL": | |
| return gr.update(visible=True), gr.update(visible=False) | |
| else: | |
| return gr.update(visible=False), gr.update(visible=True) | |
| def update_preview(url): | |
| if url: | |
| return getImg(url) | |
| else: | |
| return None | |
| with gr.Blocks(theme=gr.themes.Base(), title="Virtual Interior Staging") as demo: | |
| with gr.Column(): | |
| gr.Markdown( | |
| """ | |
| <h1 style="text-align: center;">Virtual Interior Staging</h1> | |
| <p style="text-align: center;">Easily visualize different interior design styles for your rooms.<br>Email: [email protected]</p> | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| mode = gr.Radio(["URL", "Upload Image"], label="Input Mode", value="URL", interactive=True) | |
| url_input_row = gr.Column(visible=True) | |
| with url_input_row: | |
| url = gr.Textbox(label="Image URL") | |
| url_preview = gr.Image(label="Image Preview") | |
| upload_input_row = gr.Column(visible=False) | |
| with upload_input_row: | |
| upload_image = gr.Image(label="Upload Image", type="pil") | |
| mode.change(fn=update_visibility, inputs=mode, outputs=[url_input_row, upload_input_row]) | |
| room_type = gr.Dropdown(choices=["Bedroom", "Living Room", "Kitchen", "Bathroom", "Home Office", "Dining Room"], label="Room Type") | |
| furniture_style = gr.Dropdown(choices=["Standard", "Modern", "Scandinavian", "Industrial", "Midcentury", "Luxury", "Costal", "Farmhouse"], label="Furniture Style") | |
| result_button = gr.Button("Generate Design") | |
| with gr.Column(): | |
| result_image = gr.Image(label="Staged Image") | |
| url.change(fn=update_preview, inputs=url, outputs=url_preview) | |
| result_button.click(process, inputs=[url, room_type, furniture_style, mode], outputs=result_image) | |
| demo.launch() |