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 main UI interface with all components""" # Store current variations at UI level current_variations = [] with gr.Blocks(css=""" /* Base styles */ :root { --section-title-size: 1.2rem; --section-spacing: var(--spacing-lg); --panel-min-height: auto; } /* Row styling for equal heights */ .gr-row { margin-bottom: var(--spacing-md); display: flex; align-items: stretch; } .gr-row > .gr-column { display: flex; flex-direction: column; } .gr-row > .gr-column > .gr-group { flex: 1; display: flex; flex-direction: column; } /* Consistent title styling */ .gr-markdown h2 { font-size: var(--section-title-size) !important; font-weight: 600 !important; margin: var(--spacing-sm) 0 var(--spacing-md) !important; padding: 0 !important; color: var(--body-text-color) !important; } .gr-markdown h3 { font-size: 1rem !important; font-weight: 500 !important; margin: var(--spacing-sm) 0 !important; padding: 0 !important; color: var(--body-text-color) !important; } /* Panel styling */ .gr-group { border: 1px solid var(--border-color-primary); border-radius: var(--radius-lg); padding: var(--section-spacing); margin: var(--spacing-sm) 0; background: var(--background-fill-primary); min-height: var(--panel-min-height); height: auto !important; display: flex; flex-direction: column; } /* Form and input spacing */ .gr-form { gap: var(--spacing-sm); flex-grow: 1; display: flex; flex-direction: column; } .gr-form > div { gap: var(--spacing-sm); } /* Dropdown styling */ .gr-dropdown { margin-bottom: var(--spacing-sm); } /* Gallery improvements */ #gallery { margin-top: 0; height: 300px !important; } #gallery img { object-fit: contain !important; width: 100% !important; height: 100% !important; max-height: none !important; } /* Button styling */ .button-row { display: flex; justify-content: center; padding: var(--spacing-xl) 0; } .button-row button { min-width: 200px; font-size: 1.1em; font-weight: 600; } /* Text areas */ .gr-textarea { font-family: monospace; line-height: 1.4; } .gr-textarea:disabled { opacity: 0.9; background-color: var(--background-fill-secondary); } /* Progress indicator */ .progress-bar { margin: var(--spacing-sm) 0; } /* Make dropdowns always visible */ .gr-dropdown { display: block !important; visibility: visible !important; } /* Upload area */ .upload-group { height: 100%; } .upload-group .gr-image { min-height: 300px; } /* Checkbox alignment */ .gr-checkbox-row { display: flex !important; align-items: center !important; min-height: 4.5rem !important; } .gr-checkbox-row .gr-checkbox { margin: auto 0 !important; } /* Remove any fixed heights from groups */ .gr-group > div { height: auto !important; min-height: unset !important; } /* Ensure consistent spacing in all panels */ .gr-group > div:not(:last-child) { margin-bottom: var(--spacing-sm); } /* Make surface finishes more compact */ .surface-finishes { padding: var(--spacing-xs) !important; margin: 0 !important; border: 1px solid var(--border-color-primary) !important; border-radius: var(--radius-lg) !important; background: var(--background-fill-primary) !important; } .surface-finishes .gr-form { gap: var(--spacing-xs) !important; margin: 0 !important; padding: 0 !important; flex-grow: 0 !important; } .surface-finishes .gr-dropdown { margin: 0 0 var(--spacing-xs) 0 !important; } .surface-finishes .gr-row { margin: 0 !important; gap: var(--spacing-sm) !important; } .surface-finishes .gr-group { padding: 0 !important; margin: 0 !important; border: none !important; background: none !important; box-shadow: none !important; min-height: 0 !important; } .surface-finishes .gr-markdown { margin: 0 0 var(--spacing-xs) 0 !important; } .surface-finishes .gr-form > div { gap: var(--spacing-xs) !important; margin: 0 !important; } .surface-finishes .gr-column { flex-grow: 0 !important; padding: 0 !important; } /* Remove any minimum heights */ .surface-finishes .gr-group > div { min-height: 0 !important; height: auto !important; margin: 0 !important; } /* Override any flex growth */ .surface-finishes .gr-form, .surface-finishes .gr-column, .surface-finishes .gr-group { flex: 0 0 auto !important; } /* Wall decorations and special requests */ .wall-decorations-row > .gr-column > .gr-group { height: 100%; } /* Upload and gallery row */ .upload-gallery-row > .gr-column > .gr-group { height: 100%; } """) as interface: gr.Markdown("### Interior Design Assistant") with gr.Blocks(): # Row 1 - Basic Settings with gr.Row(): with gr.Group(): gr.Markdown("## 🏠 Basic Settings") with gr.Row(): room_type = gr.Dropdown( choices=[ "Living Room", "Bedroom", "Kitchen", "Dining Room", "Bathroom", "Home Office", "Kids Room", "Master Bedroom", "Guest Room", "Studio Apartment", "Entryway", "Hallway", "Game Room", "Library", "Home Theater", "Gym" ], label="Room Type", value="Living Room" ) style_preset = gr.Dropdown( choices=[ "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="Modern" ) color_scheme = gr.Dropdown( choices=[ "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="Neutral" ) # 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=7.5, step=0.1, label="Design Freedom" ) strength = gr.Slider( minimum=0.1, maximum=1.0, value=0.75, 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") # Progress indicator progress = gr.Progress(track_tqdm=True) def upload_to_drive(image, folder_id): """Upload an image to Google Drive folder""" try: # OAuth 2.0 scopes SCOPES = ['https://www.googleapis.com/auth/drive.file'] # Start OAuth 2.0 flow flow = InstalledAppFlow.from_client_secrets_file( 'credentials.json', SCOPES ) creds = flow.run_local_server(port=0) # Build the Drive API service service = build('drive', 'v3', credentials=creds) # Convert numpy array to bytes img = Image.fromarray(image) img_byte_arr = BytesIO() img.save(img_byte_arr, format='PNG') img_byte_arr.seek(0) # Prepare the file metadata timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") file_metadata = { 'name': f'design_variation_{timestamp}.png', 'parents': [folder_id] } # Create media media = MediaIoBaseUpload( img_byte_arr, mimetype='image/png', resumable=True ) # Execute the upload file = service.files().create( body=file_metadata, media_body=media, fields='id' ).execute() print(f"File uploaded successfully. File ID: {file.get('id')}") return True except Exception as e: print(f"Error uploading to Drive: {e}") return False def extract_folder_id(drive_url): """Extract folder ID from Google Drive URL""" try: # Handle different URL formats if 'folders/' in drive_url: folder_id = drive_url.split('folders/')[-1].split('?')[0] return folder_id except Exception as e: print(f"Error extracting folder ID: {e}") return None 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: nonlocal current_variations # 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, num_variations=max(1, int(num_outputs)), prompt=prompt, num_steps=int(num_steps), guidance_scale=float(guidance_scale), strength=float(strength), seed=int(seed) if seed != -1 else None ) # Store variations current_variations = variations # 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}") current_variations = [] 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] ) # Gallery click handler def on_select(evt): nonlocal current_variations try: if isinstance(evt, list) and len(evt) > 0: # Get the clicked file path clicked_path = evt[0][0] if isinstance(evt[0], tuple) else evt[0] # Get all file paths from the gallery gallery_paths = [] for item in evt: path = item[0] if isinstance(item, tuple) else item gallery_paths.append(path) # Find which image was clicked by comparing paths selected_index = gallery_paths.index(clicked_path) if 0 <= selected_index < len(current_variations): return current_variations[selected_index] except Exception as e: print(f"Gallery selection error: {e}") return None gallery.select( fn=on_select, inputs=gallery, outputs=[] ) return interface 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): # Start with basic room and style prompt_parts = [f"Design a {style} {room.lower()} with a {colors} color scheme"] # Add floor details only if type is not "Keep Existing" if floor_t and floor_t != "Keep Existing": floor_desc = floor_t if floor_c and floor_c != "Keep Existing": floor_desc += f" in {floor_c}" if floor_p and floor_p != "Keep Existing": floor_desc += f" with {floor_p} pattern" prompt_parts.append(f"featuring {floor_desc} flooring") # Add wall details only if type is not "Keep Existing" if wall_t and wall_t != "Keep Existing": wall_desc = wall_t if wall_c and wall_c != "Keep Existing": wall_desc += f" in {wall_c}" if wall_f and wall_f != "Keep Existing": wall_desc += f" with {wall_f} finish" prompt_parts.append(f"with {wall_desc} walls") # Add accessories only if enabled AND properties are selected and not "Keep Existing" or "None" accessories = [] # Art Print if art_en and art_col and art_col not in ["Keep Existing", "None"]: accessories.append(f"{art_size} {art_col} Art Print") # Mirror if mirror_en and mirror_fr and mirror_fr not in ["Keep Existing", "None"]: accessories.append(f"{mirror_size} Mirror with {mirror_fr} frame") # Wall Sconce if sconce_en and sconce_col and sconce_col not in ["Keep Existing", "None"]: accessories.append(f"{sconce_style} {sconce_col} Wall Sconce") # Floating Shelves if shelf_en and shelf_col and shelf_col not in ["Keep Existing", "None"]: accessories.append(f"{shelf_size} {shelf_col} Floating Shelves") # Wall Plants if plants_en and plants_type and plants_type not in ["Keep Existing", "None"]: accessories.append(f"{plants_size} {plants_type}") # Only add accessories section if there are any accessories if accessories: prompt_parts.append("decorated with " + ", ".join(accessories)) # Add custom text only if provided and non-empty if custom_text and custom_text.strip(): prompt_parts.append(custom_text.strip()) return ", ".join(prompt_parts) 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()