import gradio as gr from PIL import Image import numpy as np import time import os import random from typing import List import traceback from google.oauth2.credentials import Credentials from google_auth_oauthlib.flow import InstalledAppFlow from googleapiclient.discovery import build from googleapiclient.http import MediaIoBaseUpload from io import BytesIO import datetime # Import the model interface from model import DesignModel # For testing, import the mock model from mock_model import MockDesignModel def create_ui(model: DesignModel): """Create the user interface for the application""" with gr.Blocks() as interface: gr.Markdown("## 🏠 Basic Settings") with gr.Row(): with gr.Column(): room_type = gr.Dropdown( choices=[ "None", "Living Room", "Bedroom", "Dining Room", "Kitchen", "Bathroom", "Home Office", "Master Bedroom", "Guest Room", "Family Room", "Study", "Game Room", "Library", "Nursery", "Craft Room" ], label="Room Type", value="None" ) style_preset = gr.Dropdown( choices=[ "None", "Modern", "Contemporary", "Minimalist", "Industrial", "Scandinavian", "Mid-Century Modern", "Traditional", "Transitional", "Farmhouse", "Rustic", "Bohemian", "Art Deco", "Coastal", "Mediterranean", "Japanese", "French Country", "Victorian", "Colonial", "Gothic", "Baroque", "Rococo", "Neoclassical", "Eclectic", "Zen", "Tropical", "Shabby Chic", "Hollywood Regency", "Southwestern", "Asian Fusion", "Retro" ], label="Design Style", value="None" ) color_scheme = gr.Dropdown( choices=[ "None", "Neutral", "Monochromatic", "Minimalist White", "Warm Gray", "Cool Gray", "Earth Tones", "Pastel", "Bold Primary", "Jewel Tones", "Black and White", "Navy and Gold", "Forest Green", "Desert Sand", "Ocean Blue", "Sunset Orange", "Deep Purple", "Emerald Green", "Ruby Red", "Sapphire Blue", "Golden Yellow", "Sage Green", "Dusty Rose", "Charcoal", "Cream", "Burgundy", "Teal", "Copper", "Silver", "Bronze", "Slate" ], label="Color Mood", value="None" ) # Row 2 - Surface Finishes with gr.Row(): # Floor Options with gr.Column(scale=1): with gr.Group(): gr.Markdown("## 🎨 Floor Options") floor_type = gr.Dropdown( choices=[ "Keep Existing", "Hardwood", "Stone Tiles", "Porcelain Tiles", "Soft Carpet", "Polished Concrete", "Marble", "Vinyl", "Natural Bamboo", "Cork", "Ceramic Tiles", "Terrazzo", "Slate", "Travertine", "Laminate", "Engineered Wood", "Mosaic Tiles", "Luxury Vinyl Tiles", "Stained Concrete" ], label="Material", value="Keep Existing" ) floor_color = gr.Dropdown( choices=[ "Keep Existing", "Light Oak", "Rich Walnut", "Cool Gray", "Whitewashed", "Warm Cherry", "Deep Brown", "Classic Black", "Natural", "Sandy Beige", "Chocolate", "Espresso", "Honey Oak", "Weathered Gray", "White Marble", "Cream Travertine", "Dark Slate", "Golden Teak", "Rustic Pine", "Ebony" ], label="Color", value="Keep Existing" ) floor_pattern = gr.Dropdown( choices=[ "Keep Existing", "Classic Straight", "Elegant Herringbone", "V-Pattern", "Decorative Parquet", "Diagonal Layout", "Basketweave", "Chevron", "Random Length", "Grid Pattern", "Versailles Pattern", "Running Bond", "Hexagonal", "Moroccan Pattern", "Brick Layout", "Diamond Pattern", "Windmill Pattern", "Large Format", "Mixed Width" ], label="Pattern", value="Keep Existing" ) # Wall Options with gr.Column(scale=1): with gr.Group(): gr.Markdown("## 🎨 Wall Options") wall_type = gr.Dropdown( choices=[ "Keep Existing", "Fresh Paint", "Designer Wallpaper", "Textured Finish", "Wood Panels", "Exposed Brick", "Natural Stone", "Wooden Planks", "Modern Concrete", "Venetian Plaster", "Wainscoting", "Shiplap", "3D Wall Panels", "Fabric Panels", "Metal Panels", "Cork Wall", "Tile Feature", "Glass Panels", "Acoustic Panels", "Living Wall" ], label="Treatment", value="Keep Existing" ) wall_color = gr.Dropdown( choices=[ "Keep Existing", "Crisp White", "Soft White", "Warm Beige", "Gentle Gray", "Sky Blue", "Nature Green", "Sunny Yellow", "Blush Pink", "Deep Blue", "Bold Black", "Sage Green", "Terracotta", "Navy Blue", "Charcoal Gray", "Lavender", "Olive Green", "Dusty Rose", "Teal", "Burgundy" ], label="Color", value="Keep Existing" ) wall_finish = gr.Dropdown( choices=[ "Keep Existing", "Soft Matte", "Subtle Eggshell", "Pearl Satin", "Sleek Semi-Gloss", "High Gloss", "Suede Texture", "Metallic", "Chalk Finish", "Distressed", "Brushed", "Smooth", "Textured", "Venetian", "Lime Wash", "Concrete", "Rustic", "Lacquered", "Hammered", "Patina" ], label="Finish", value="Keep Existing" ) # Row 3 - Wall Decorations and Special Requests with gr.Row(elem_classes="wall-decorations-row"): # Wall Decorations with gr.Column(scale=2): with gr.Group(): gr.Markdown("## 🖼️ Wall Decorations") # Art and Mirror with gr.Row(): # Art Print with gr.Column(): with gr.Row(): art_print_enable = gr.Checkbox(label="Add Artwork", value=False) art_print_color = gr.Dropdown( choices=[ "None", "Classic Black & White", "Vibrant Colors", "Single Color", "Soft Colors", "Modern Abstract", "Earth Tones", "Pastel Palette", "Bold Primary Colors", "Metallic Accents", "Monochromatic", "Jewel Tones", "Watercolor", "Vintage Colors", "Neon Accents", "Natural Hues", "Ocean Colors", "Desert Palette" ], label="Art Style", value="None" ) art_print_size = gr.Dropdown( choices=[ "None", "Modest", "Standard", "Statement", "Oversized", "Gallery Wall", "Diptych", "Triptych", "Mini Series", "Floor to Ceiling", "Custom Size" ], label="Art Size", value="None" ) # Mirror with gr.Column(): with gr.Row(): mirror_enable = gr.Checkbox(label="Add Mirror", value=False) mirror_frame = gr.Dropdown( choices=[ "None", "Gold", "Silver", "Black", "White", "Wood", "Brass", "Bronze", "Copper", "Chrome", "Antique Gold", "Brushed Nickel", "Rustic Wood", "Ornate", "Minimalist", "LED Backlit", "Bamboo", "Rattan", "Leather Wrapped" ], label="Frame Style", value="None" ) mirror_size = gr.Dropdown( choices=[ "Small", "Medium", "Large", "Full Length", "Oversized", "Double Width", "Floor Mirror", "Vanity Size", "Statement Piece", "Custom Size" ], label="Mirror Size", value="Medium" ) # Sconce, Shelf, and Plants with gr.Row(): # Sconce with gr.Column(): with gr.Row(): sconce_enable = gr.Checkbox(label="Add Wall Sconce", value=False) sconce_color = gr.Dropdown( choices=[ "None", "Black", "Gold", "Silver", "Bronze", "White", "Brass", "Copper", "Chrome", "Antique Brass", "Brushed Nickel", "Oil-Rubbed Bronze", "Pewter", "Rose Gold", "Matte Black", "Polished Nickel", "Aged Brass", "Champagne", "Gunmetal" ], label="Sconce Color", value="None" ) sconce_style = gr.Dropdown( choices=[ "Modern", "Traditional", "Industrial", "Art Deco", "Minimalist", "Vintage", "Contemporary", "Rustic", "Coastal", "Farmhouse", "Mid-Century", "Bohemian", "Scandinavian", "Asian", "Mediterranean", "Gothic", "Transitional", "Eclectic", "Victorian" ], label="Sconce Style", value="Modern" ) # Floating Shelves with gr.Column(): with gr.Row(): shelf_enable = gr.Checkbox(label="Add Floating Shelves", value=False) shelf_color = gr.Dropdown( choices=[ "None", "White", "Black", "Natural Wood", "Glass", "Dark Wood", "Light Wood", "Metal", "Gold", "Silver", "Bronze", "Reclaimed Wood", "Bamboo", "Marble", "Industrial Metal", "Two-Tone", "Concrete", "Acrylic", "Copper", "Brass" ], label="Shelf Material", value="None" ) shelf_size = gr.Dropdown( choices=[ "Small", "Medium", "Large", "Set of 3", "Extra Long", "Corner Set", "Asymmetric Set", "Graduated Sizes", "Custom Length", "Mini Cubes", "Full Wall", "Mixed Sizes", "Modular System" ], label="Shelf Size", value="Medium" ) # Plants with gr.Column(): with gr.Row(): plants_enable = gr.Checkbox(label="Add Plants", value=False) plants_type = gr.Dropdown( choices=[ "None", "Hanging Plants", "Vertical Garden", "Plant Shelf", "Single Plant", "Climbing Vines", "Air Plants", "Succulent Wall", "Herb Garden", "Mixed Tropical", "Fern Collection", "Living Wall", "Moss Wall", "Potted Arrangement", "Plant Corner", "Cascading Plants", "Bamboo Screen", "Terrarium Wall" ], label="Plant Type", value="None" ) plants_size = gr.Dropdown( choices=[ "Small", "Medium", "Large", "Mixed Sizes", "Full Wall", "Statement Piece", "Compact", "Expansive", "Accent", "Floor to Ceiling", "Window Height", "Custom Size", "Modular" ], label="Plant Coverage", value="Medium" ) # Special Requests and Advanced Settings with gr.Column(scale=1): with gr.Group(): gr.Markdown("## ✨ Special Requests") input_text = gr.Textbox( label="Additional Details", placeholder="Add any special requests or details here...", lines=3 ) num_outputs = gr.Slider( minimum=1, maximum=50, value=1, step=1, label="Number of Variations" ) gr.Markdown("### Advanced Settings") num_steps = gr.Slider( minimum=20, maximum=100, value=50, step=1, label="Quality Steps" ) guidance_scale = gr.Slider( minimum=1, maximum=20, value=10.0, step=0.1, label="Design Freedom" ) strength = gr.Slider( minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Change Amount" ) seed = gr.Number( label="Seed (leave empty for random)", value=-1, precision=0 ) with gr.Row(): save_to_drive = gr.Checkbox(label="Save to Google Drive") drive_url = gr.Textbox( label="Drive Folder URL", placeholder="https://drive.google.com/drive/folders/..." ) # Row 4 - Current Prompts with gr.Row(): with gr.Group(): gr.Markdown("## 📝 Current Prompts") prompt_display = gr.TextArea( label="Positive Prompt", interactive=False, lines=3, value="Your design prompt will appear here..." ) negative_prompt = gr.TextArea( label="Negative Prompt", value="blurry, low quality, distorted, deformed, disfigured, watermark, text, bad proportions, duplicate, double, multiple, broken, cropped", lines=2, interactive=False ) # Row 5 - Upload and Gallery with gr.Row(elem_classes="upload-gallery-row"): # Upload Area with gr.Column(scale=1): with gr.Group(): gr.Markdown("## 📸 Upload Photo") input_image = gr.Image( label="Upload a photo of your room", type='pil' ) # Gallery Area with gr.Column(scale=2): with gr.Group(): gr.Markdown("## 🖼️ Generated Variations") gallery = gr.Gallery( show_label=False, elem_id="gallery", columns=4, rows=1, height="300px", object_fit="contain", preview=True, show_share_button=False ) # Row 6 - Create Button with gr.Row(elem_classes="button-row"): submit = gr.Button("✨ Create My Design", variant="primary", size="lg") def is_valid_value(*values): """Helper function to check if values are valid (not None or Keep Existing)""" return all(v not in ["None", "Keep Existing"] for v in values) def update_prompt(room, style, colors, floor_t, floor_c, floor_p, wall_t, wall_c, wall_f, custom_text, art_en, art_col, art_size, mirror_en, mirror_fr, mirror_size, sconce_en, sconce_col, sconce_style, shelf_en, shelf_col, shelf_size, plants_en, plants_type, plants_size): """Generate a prompt for the design, skipping any None or Keep Existing values.""" invalid_values = {"None", "Keep Existing", None, ""} prompt_parts = [] # Basic room settings - each part separate for maximum flexibility if all(x not in invalid_values for x in [style, room, colors]): base_parts = [] if style not in invalid_values: base_parts.append(style) if room not in invalid_values: base_parts.append(room.lower()) if colors not in invalid_values: base_parts.append(f"with a {colors} color scheme") if base_parts: prompt_parts.append("Design a " + " ".join(base_parts)) # Floor details - build up piece by piece if floor_t not in invalid_values: floor_parts = [floor_t] if floor_c not in invalid_values: floor_parts.append(f"in {floor_c}") if floor_p not in invalid_values: floor_parts.append(f"with {floor_p} pattern") prompt_parts.append(f"featuring {' '.join(floor_parts)} flooring") # Wall details - build up piece by piece if wall_t not in invalid_values: wall_parts = [wall_t] if wall_c not in invalid_values: wall_parts.append(f"in {wall_c}") if wall_f not in invalid_values: wall_parts.append(f"with {wall_f} finish") prompt_parts.append(f"with {' '.join(wall_parts)} walls") # Accessories - only add valid combinations accessories = [] # Art Print if art_en and art_col not in invalid_values and art_size not in invalid_values: accessories.append(f"{art_size} {art_col} Art Print") # Mirror if mirror_en and mirror_fr not in invalid_values and mirror_size not in invalid_values: accessories.append(f"{mirror_size} Mirror with {mirror_fr} frame") # Wall Sconce if sconce_en and sconce_col not in invalid_values and sconce_style not in invalid_values: accessories.append(f"{sconce_style} {sconce_col} Wall Sconce") # Floating Shelves if shelf_en and shelf_col not in invalid_values and shelf_size not in invalid_values: accessories.append(f"{shelf_size} {shelf_col} Floating Shelves") # Wall Plants if plants_en and plants_type not in invalid_values and plants_size not in invalid_values: accessories.append(f"{plants_size} {plants_type}") if accessories: prompt_parts.append("decorated with " + ", ".join(accessories)) # Add custom text if present if custom_text and custom_text.strip() and custom_text not in invalid_values: prompt_parts.append(custom_text.strip()) return ", ".join(prompt_parts) if prompt_parts else "Please select room type, style, and color scheme" def on_submit(image, room, style, colors, floor_t, floor_c, floor_p, wall_t, wall_c, wall_f, custom_text, art_en, art_col, art_size, mirror_en, mirror_fr, mirror_size, sconce_en, sconce_col, sconce_style, shelf_en, shelf_col, shelf_size, plants_en, plants_type, plants_size, num_outputs, save_to_drive, drive_url, num_steps, guidance_scale, seed, strength): if image is None: return [] try: # Generate the prompt prompt = update_prompt( room, style, colors, floor_t, floor_c, floor_p, wall_t, wall_c, wall_f, custom_text, art_en, art_col, art_size, mirror_en, mirror_fr, mirror_size, sconce_en, sconce_col, sconce_style, shelf_en, shelf_col, shelf_size, plants_en, plants_type, plants_size ) # Generate variations variations = model.generate_design( image=image, prompt=prompt, num_variations=max(1, int(num_outputs)), num_steps=int(num_steps), guidance_scale=float(guidance_scale), strength=float(strength), seed=int(seed) if seed != -1 else None ) # Handle Google Drive upload if enabled if save_to_drive and drive_url: folder_id = extract_folder_id(drive_url) if folder_id: for variation in variations: upload_to_drive(variation, folder_id) # Convert variations to gallery format gallery_images = [(v, None) for v in variations] return gallery_images except Exception as e: print(f"Error in generation: {e}") return [] submit.click( on_submit, inputs=[ input_image, room_type, style_preset, color_scheme, floor_type, floor_color, floor_pattern, wall_type, wall_color, wall_finish, input_text, art_print_enable, art_print_color, art_print_size, mirror_enable, mirror_frame, mirror_size, sconce_enable, sconce_color, sconce_style, shelf_enable, shelf_color, shelf_size, plants_enable, plants_type, plants_size, num_outputs, save_to_drive, drive_url, num_steps, guidance_scale, seed, strength ], outputs=[gallery] ) # Update prompt display when any input changes def update_prompt_display(*args): try: prompt = update_prompt(*args) return [prompt, negative_prompt.value] # Return both prompts except Exception as e: print(f"Error updating prompt: {e}") return ["Error generating prompt", negative_prompt.value] # List of all inputs that should trigger prompt updates prompt_inputs = [ room_type, style_preset, color_scheme, floor_type, floor_color, floor_pattern, wall_type, wall_color, wall_finish, input_text, art_print_enable, art_print_color, art_print_size, mirror_enable, mirror_frame, mirror_size, sconce_enable, sconce_color, sconce_style, shelf_enable, shelf_color, shelf_size, plants_enable, plants_type, plants_size ] # Connect all inputs to prompt update for input_component in prompt_inputs: input_component.change( fn=update_prompt_display, inputs=prompt_inputs, outputs=[prompt_display, negative_prompt] ) return interface def main(): """Main entry point for the application""" import sys # Check if we're in test mode is_test_mode = "--test" in sys.argv if is_test_mode: print("Starting in TEST mode...") from mock_model import MockDesignModel model = MockDesignModel() else: print("Starting in PRODUCTION mode...") from prod_model import ProductionDesignModel model = ProductionDesignModel() interface = create_ui(model) interface.launch( share=False, show_api=False, # Hide API docs show_error=True # Show errors for debugging ) if __name__ == "__main__": main()