StableDesign2

Runtime error

App Files Files Community

anbucur commited on Jan 19

Commit

4e4b650

1 Parent(s): 39751c2

Added functionality

Browse files

Files changed (8) hide show

README.md +107 -0
app.py +888 -364
credentials.json +1 -0
mock_model.py +83 -0
model.py +14 -0
prod_model.py +170 -0
requirements.txt +33 -10
test_prompt.py +320 -0

README.md CHANGED Viewed

@@ -11,3 +11,110 @@ license: mit
 ---
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

 ---
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Interior Design Assistant
+An AI-powered interior design assistant that helps transform room photos with customizable design variations.
+## Features
+- Upload room photos and generate design variations
+- Customize room type, style, color schemes
+- Adjust floor and wall finishes
+- Add wall decorations (art, mirrors, sconces)
+- Control generation parameters (quality, creativity, etc.)
+- Save results to Google Drive
+- Production and test modes
+## Requirements
+- Python 3.8 or higher
+- CUDA-capable GPU with 8GB+ VRAM (recommended)
+- CPU-only mode supported but slower
+## Installation
+1. Clone the repository:
+```bash
+git clone [your-repo-url]
+cd StableDesign2
+```
+2. Create and activate a virtual environment:
+```bash
+python -m venv venv
+source venv/bin/activate  # Linux/Mac
+# or
+.\venv\Scripts\activate  # Windows
+```
+3. Install dependencies:
+```bash
+pip install -r requirements.txt
+```
+4. Set up Google Drive integration (optional):
+- Create a project in Google Cloud Console
+- Enable the Google Drive API
+- Create OAuth 2.0 credentials
+- Download credentials and save as `credentials.json` in the project root
+## Usage
+### Production Mode
+```bash
+python app.py
+```
+### Test Mode (for development)
+```bash
+python app.py --test
+```
+The interface will be available at `http://localhost:7860`
+## Configuration
+### Model Settings
+- Quality Steps: 20-100 (default: 50)
+- Design Freedom: 1-20 (default: 7.5)
+- Change Amount: 0.1-1.0 (default: 0.75)
+- Number of Variations: 1-4
+### Design Options
+- Room Types: 16 options
+- Design Styles: 20 options
+- Color Schemes: 20 options
+- Floor & Wall Options: Multiple materials, colors, patterns
+- Wall Decorations: Art, mirrors, sconces, shelves, plants
+## Error Handling
+The application includes comprehensive error handling:
+- Input validation
+- Model generation fallbacks
+- Google Drive upload retries
+- Detailed error logging
+## Development
+### Running Tests
+```bash
+python -m pytest
+```
+### Code Style
+```bash
+black .
+flake8
+isort .
+```
+## License
+[Your License]
+## Credits
+- Built with [Gradio](https://gradio.app/)
+- Powered by [Stable Diffusion](https://stability.ai/)

app.py CHANGED Viewed

@@ -1,377 +1,901 @@
-import spaces
-from typing import Tuple, Union, List
-import os
-import numpy as np
-from PIL import Image
-import torch
-from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline
-from diffusers import ControlNetModel, UniPCMultistepScheduler, AutoPipelineForText2Image
-from transformers import AutoImageProcessor, UperNetForSemanticSegmentation, AutoModelForDepthEstimation
-from colors import ade_palette
-from utils import map_colors_rgb
-from diffusers import StableDiffusionXLPipeline
 import gradio as gr
-import gc
-device = "cuda"
-dtype = torch.float16
-css = """
-#img-display-container {
-    max-height: 50vh;
-    }
-#img-display-input {
-    max-height: 40vh;
-    }
-#img-display-output {
-    max-height: 40vh;
-    }
-"""
-def filter_items(
-    colors_list: Union[List, np.ndarray],
-    items_list: Union[List, np.ndarray],
-    items_to_remove: Union[List, np.ndarray]
-) -> Tuple[Union[List, np.ndarray], Union[List, np.ndarray]]:
-    """
-    Filters items and their corresponding colors from given lists, excluding
-    specified items.
-    Args:
-        colors_list: A list or numpy array of colors corresponding to items.
-        items_list: A list or numpy array of items.
-        items_to_remove: A list or numpy array of items to be removed.
-    Returns:
-        A tuple of two lists or numpy arrays: filtered colors and filtered
-        items.
-    """
-    filtered_colors = []
-    filtered_items = []
-    for color, item in zip(colors_list, items_list):
-        if item not in items_to_remove:
-            filtered_colors.append(color)
-            filtered_items.append(item)
-    return filtered_colors, filtered_items
-def get_segmentation_pipeline(
-) -> Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]:
-    """Method to load the segmentation pipeline
-    Returns:
-        Tuple[AutoImageProcessor, UperNetForSemanticSegmentation]: segmentation pipeline
-    """
-    image_processor = AutoImageProcessor.from_pretrained(
-        "openmmlab/upernet-convnext-small"
-    )
-    image_segmentor = UperNetForSemanticSegmentation.from_pretrained(
-        "openmmlab/upernet-convnext-small"
-    )
-    return image_processor, image_segmentor
-@torch.inference_mode()
-@spaces.GPU
-def segment_image(
-        image: Image,
-        image_processor: AutoImageProcessor,
-        image_segmentor: UperNetForSemanticSegmentation
-) -> Image:
-    """
-    Segments an image using a semantic segmentation model.
-    Args:
-        image (Image): The input image to be segmented.
-        image_processor (AutoImageProcessor): The processor to prepare the
-            image for segmentation.
-        image_segmentor (UperNetForSemanticSegmentation): The semantic
-            segmentation model used to identify different segments in the image.
-    Returns:
-        Image: The segmented image with each segment colored differently based
-            on its identified class.
-    """
-    # image_processor, image_segmentor = get_segmentation_pipeline()
-    pixel_values = image_processor(image, return_tensors="pt").pixel_values
-    with torch.no_grad():
-        outputs = image_segmentor(pixel_values)
-    seg = image_processor.post_process_semantic_segmentation(
-        outputs, target_sizes=[image.size[::-1]])[0]
-    color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
-    palette = np.array(ade_palette())
-    for label, color in enumerate(palette):
-        color_seg[seg == label, :] = color
-    color_seg = color_seg.astype(np.uint8)
-    seg_image = Image.fromarray(color_seg).convert('RGB')
-    return seg_image
-def get_depth_pipeline():
-    feature_extractor = AutoImageProcessor.from_pretrained("LiheYoung/depth-anything-large-hf",
-                                                           torch_dtype=dtype)
-    depth_estimator = AutoModelForDepthEstimation.from_pretrained("LiheYoung/depth-anything-large-hf",
-                                                                  torch_dtype=dtype)
-    return feature_extractor, depth_estimator
-@torch.inference_mode()
-@spaces.GPU
-def get_depth_image(
-        image: Image,
-        feature_extractor: AutoImageProcessor,
-        depth_estimator: AutoModelForDepthEstimation
-) -> Image:
-    image_to_depth = feature_extractor(images=image, return_tensors="pt").to(device)
-    with torch.no_grad():
-        depth_map = depth_estimator(**image_to_depth).predicted_depth
-    width, height = image.size
-    depth_map = torch.nn.functional.interpolate(
-        depth_map.unsqueeze(1).float(),
-        size=(height, width),
-        mode="bicubic",
-        align_corners=False,
-    )
-    depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
-    depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
-    depth_map = (depth_map - depth_min) / (depth_max - depth_min)
-    image = torch.cat([depth_map] * 3, dim=1)
-    image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
-    image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
-    return image
-def resize_dimensions(dimensions, target_size):
-    """
-    Resize PIL to target size while maintaining aspect ratio
-    If smaller than target size leave it as is
-    """
-    width, height = dimensions
-    # Check if both dimensions are smaller than the target size
-    if width < target_size and height < target_size:
-        return dimensions
-    # Determine the larger side
-    if width > height:
-        # Calculate the aspect ratio
-        aspect_ratio = height / width
-        # Resize dimensions
-        return (target_size, int(target_size * aspect_ratio))
-    else:
-        # Calculate the aspect ratio
-        aspect_ratio = width / height
-        # Resize dimensions
-        return (int(target_size * aspect_ratio), target_size)
-def flush():
-    gc.collect()
-    torch.cuda.empty_cache()
-class ControlNetDepthDesignModelMulti:
-    """ Produces random noise images """
-    def __init__(self):
-        """ Initialize your model(s) here """
-        #os.environ['HF_HUB_OFFLINE'] = "True"
-        self.seed = 323*111
-        self.neg_prompt = "window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner"
-        self.control_items = ["windowpane;window", "door;double;door"]
-        self.additional_quality_suffix = "interior design, 4K, high resolution, photorealistic"
-    @spaces.GPU
-    def generate_design(self, empty_room_image: Image, prompt: str, guidance_scale: int = 10, num_steps: int = 50, strength: float =0.9, img_size: int = 640) -> Image:
-        """
-        Given an image of an empty room and a prompt
-        generate the designed room according to the prompt
-        Inputs -
-            empty_room_image - An RGB PIL Image of the empty room
-            prompt - Text describing the target design elements of the room
-        Returns -
-            design_image - PIL Image of the same size as the empty room image
-                           If the size is not the same the submission will fail.
-        """
-        print(prompt)
-        flush()
-        self.generator = torch.Generator(device=device).manual_seed(self.seed)
-        pos_prompt = prompt + f', {self.additional_quality_suffix}'
-        orig_w, orig_h = empty_room_image.size
-        new_width, new_height = resize_dimensions(empty_room_image.size, img_size)
-        input_image = empty_room_image.resize((new_width, new_height))
-        real_seg = np.array(segment_image(input_image,
-                                          seg_image_processor,
-                                          image_segmentor))
-        unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
-        unique_colors = [tuple(color) for color in unique_colors]
-        segment_items = [map_colors_rgb(i) for i in unique_colors]
-        chosen_colors, segment_items = filter_items(
-            colors_list=unique_colors,
-            items_list=segment_items,
-            items_to_remove=self.control_items
         )
-        mask = np.zeros_like(real_seg)
-        for color in chosen_colors:
-            color_matches = (real_seg == color).all(axis=2)
-            mask[color_matches] = 1
-        image_np = np.array(input_image)
-        image = Image.fromarray(image_np).convert("RGB")
-        mask_image = Image.fromarray((mask * 255).astype(np.uint8)).convert("RGB")
-        segmentation_cond_image = Image.fromarray(real_seg).convert("RGB")
-        image_depth = get_depth_image(image, depth_feature_extractor, depth_estimator)
-        # generate image that would be used as IP-adapter
-        flush()
-        new_width_ip = int(new_width / 8) * 8
-        new_height_ip = int(new_height / 8) * 8
-        ip_image = guide_pipe(pos_prompt,
-                                   num_inference_steps=num_steps,
-                                   negative_prompt=self.neg_prompt,
-                                   height=new_height_ip,
-                                   width=new_width_ip,
-                                   generator=[self.generator]).images[0]
-        flush()
-        generated_image = pipe(
-            prompt=pos_prompt,
-            negative_prompt=self.neg_prompt,
-            num_inference_steps=num_steps,
-            strength=strength,
-            guidance_scale=guidance_scale,
-            generator=[self.generator],
-            image=image,
-            mask_image=mask_image,
-            ip_adapter_image=ip_image,
-            control_image=[image_depth, segmentation_cond_image],
-            controlnet_conditioning_scale=[0.5, 0.5]
-        ).images[0]
-        flush()
-        design_image = generated_image.resize(
-            (orig_w, orig_h), Image.Resampling.LANCZOS
         )
-        return design_image
-def create_demo(model):
-    gr.Markdown("### Stable Design demo")
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(label="Input Image", type='pil', elem_id='img-display-input')
-            input_text = gr.Textbox(label='Prompt', placeholder='Please upload your image first', lines=2)
-            with gr.Accordion('Advanced options', open=False):
-                num_steps = gr.Slider(label='Steps',
-                                      minimum=1,
-                                      maximum=50,
-                                      value=50,
-                                      step=1)
-                img_size = gr.Slider(label='Image size',
-                                      minimum=256,
-                                      maximum=768,
-                                      value=768,
-                                      step=64)
-                guidance_scale = gr.Slider(label='Guidance Scale',
-                                           minimum=0.1,
-                                           maximum=30.0,
-                                           value=10.0,
-                                           step=0.1)
-                seed = gr.Slider(label='Seed',
-                                 minimum=-1,
-                                 maximum=2147483647,
-                                 value=323*111,
-                                 step=1,
-                                 randomize=True)
-                strength = gr.Slider(label='Strength',
-                                           minimum=0.1,
-                                           maximum=1.0,
-                                           value=0.9,
-                                           step=0.1)
-                a_prompt = gr.Textbox(
-                    label='Added Prompt',
-                    value="interior design, 4K, high resolution, photorealistic")
-                n_prompt = gr.Textbox(
-                    label='Negative Prompt',
-                    value="window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner")
-            submit = gr.Button("Submit")
-        with gr.Column():
-            design_image = gr.Image(label="Output Mask", elem_id='img-display-output')
-    def on_submit(image, text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt, img_size):
-        model.seed = seed
-        model.neg_prompt = n_prompt
-        model.additional_quality_suffix = a_prompt
-        with torch.no_grad():
-            out_img = model.generate_design(image, text, guidance_scale=guidance_scale, num_steps=num_steps, strength=strength, img_size=img_size)
-        return out_img
-    submit.click(on_submit, inputs=[input_image, input_text, num_steps, guidance_scale, seed, strength, a_prompt, n_prompt, img_size], outputs=design_image)
-    examples = gr.Examples(examples=[["imgs/bedroom_1.jpg", "An elegantly appointed bedroom in the Art Deco style, featuring a grand king-size bed with geometric bedding, a luxurious velvet armchair, and a mirrored nightstand that reflects the room's opulence. Art Deco-inspired artwork adds a touch of glamour"], ["imgs/bedroom_2.jpg", "A bedroom that exudes French country charm with a soft upholstered bed, walls adorned with floral wallpaper, and a vintage wooden wardrobe. A crystal chandelier casts a warm, inviting glow over the space"], ["imgs/dinning_room_1.jpg", "A cozy dining room that captures the essence of rustic charm with a solid wooden farmhouse table at its core, surrounded by an eclectic mix of mismatched chairs. An antique sideboard serves as a statement piece, and the ambiance is warmly lit by a series of quaint Edison bulbs dangling from the ceiling"], ["imgs/dinning_room_3.jpg", "A dining room that epitomizes contemporary elegance, anchored by a sleek, minimalist dining table paired with stylish modern chairs. Artistic lighting fixtures create a focal point above, while the surrounding minimalist decor ensures the space feels open, airy, and utterly modern"], ["imgs/image_1.jpg", "A glamorous master bedroom in Hollywood Regency style, boasting a plush tufted headboard, mirrored furniture reflecting elegance, luxurious fabrics in rich textures, and opulent gold accents for a touch of luxury."], ["imgs/image_2.jpg", "A vibrant living room with a tropical theme, complete with comfortable rattan furniture, large leafy plants bringing the outdoors in, bright cushions adding pops of color, and bamboo blinds for natural light control."], ["imgs/living_room_1.jpg", "A stylish living room embracing mid-century modern aesthetics, featuring a vintage teak coffee table at its center, complemented by a classic sunburst clock on the wall and a cozy shag rug underfoot, creating a warm and inviting atmosphere"]],
-                           inputs=[input_image, input_text], cache_examples=False)
-controlnet_depth= ControlNetModel.from_pretrained(
-    "controlnet_depth", torch_dtype=dtype, use_safetensors=True)
-controlnet_seg = ControlNetModel.from_pretrained(
-    "own_controlnet", torch_dtype=dtype, use_safetensors=True)
-pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
-    "SG161222/Realistic_Vision_V5.1_noVAE",
-    #"models/runwayml--stable-diffusion-inpainting",
-    controlnet=[controlnet_depth, controlnet_seg],
-    safety_checker=None,
-    torch_dtype=dtype
-)
-pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models",
-                     weight_name="ip-adapter_sd15.bin")
-pipe.set_ip_adapter_scale(0.4)
-pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
-pipe = pipe.to(device)
-guide_pipe = StableDiffusionXLPipeline.from_pretrained("segmind/SSD-1B",
-                                                       torch_dtype=dtype, use_safetensors=True, variant="fp16")
-guide_pipe = guide_pipe.to(device)
-seg_image_processor, image_segmentor = get_segmentation_pipeline()
-depth_feature_extractor, depth_estimator = get_depth_pipeline()
-depth_estimator = depth_estimator.to(device)
 def main():
-    model = ControlNetDepthDesignModelMulti()
-    print('Models uploaded successfully')
-    title = "# StableDesign"
-    description = """
-    <p style='font-size: 14px; margin-bottom: 10px;'><a href='https://www.linkedin.com/in/mykola-lavreniuk/'>Mykola Lavreniuk</a>, <a href='https://www.linkedin.com/in/bartosz-ludwiczuk-a677a760/'>Bartosz Ludwiczuk</a></p>
-    <p style='font-size: 16px; margin-bottom: 0px; margin-top=0px;'>Official demo for <strong>StableDesign:</strong> 2nd place solution for the Generative Interior Design 2024 <a href='https://www.aicrowd.com/challenges/generative-interior-design-challenge-2024/leaderboards?challenge_round_id=1314'>competition</a>. StableDesign is a deep learning model designed to harness the power of AI, providing innovative and creative tools for designers. Using our algorithms, images of empty rooms can be transformed into fully furnished spaces based on text descriptions. Please refer to our <a href='https://github.com/Lavreniuk/generative-interior-design'>GitHub</a> for more details.</p>
-    """
-    with gr.Blocks() as demo:
-        gr.Markdown(title)
-        gr.Markdown(description)
-        create_demo(model)
-        gr.HTML('''<br><br><br><center>You can duplicate this Space to skip the queue:<a href="https://huggingface.co/spaces/MykolaL/StableDesign?duplicate=true"><img src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a><br>
-                <p><img src="https://visitor-badge.glitch.me/badge?page_id=MykolaL/StableDesign" alt="visitors"></p></center>''')
-    demo.queue().launch(share=False)
-if __name__ == '__main__':
     main()

 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()

credentials.json ADDED Viewed

	@@ -0,0 +1 @@

+ {"installed":{"client_id":"924767856297-cdnm065phq1cq7ncfgh33bau61hjrpr8.apps.googleusercontent.com","project_id":"stabledesign","auth_uri":"https://accounts.google.com/o/oauth2/auth","token_uri":"https://oauth2.googleapis.com/token","auth_provider_x509_cert_url":"https://www.googleapis.com/oauth2/v1/certs","client_secret":"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","redirect_uris":["http://localhost"]}}

mock_model.py ADDED Viewed

	@@ -0,0 +1,83 @@

+from model import DesignModel
+from PIL import Image
+import numpy as np
+import random
+from typing import List
+import traceback
+class MockDesignModel(DesignModel):
+    def __init__(self):
+        super().__init__()
+        # Define base colors with proper alpha values
+        self.base_colors = [
+            (255, 0, 0),      # Red
+            (0, 255, 0),      # Green
+            (0, 0, 255),      # Blue
+            (255, 255, 0),    # Yellow
+            (255, 0, 255),    # Magenta
+            (0, 255, 255),    # Cyan
+            (128, 0, 0),      # Maroon
+            (0, 128, 0),      # Dark Green
+            (0, 0, 128),      # Navy
+        ]
+        # Add test-specific attributes
+        self.seed = 323*111
+        self.neg_prompt = "window, door, low resolution, banner, logo, watermark, text"
+        self.additional_quality_suffix = "interior design, 4K, high resolution"
+    def apply_tint(self, img_array: np.ndarray, color: tuple) -> np.ndarray:
+        """Apply a color tint to an image array"""
+        # Create tint array
+        tint = np.array(color, dtype=np.float32) / 255.0
+        # Apply tint with alpha blending
+        alpha = 0.3  # 30% tint strength
+        tinted = img_array * (1 - alpha) + (img_array * tint) * alpha
+        # Ensure values are within valid range
+        return np.clip(tinted, 0, 255).astype(np.uint8)
+    def generate_design(self, image: Image.Image, num_variations: int = 1, **kwargs) -> List[np.ndarray]:
+        """Generate multiple variations of the input image with different color tints"""
+        try:
+            print(f"Starting generation of {num_variations} variations")
+            # Convert image to numpy array once
+            img_array = np.array(image.convert('RGB'))
+            # Generate base colors for all variations
+            colors_needed = max(1, int(num_variations))
+            colors = []
+            # Add base colors first
+            colors.extend(self.base_colors)
+            # Generate additional random colors if needed
+            while len(colors) < colors_needed:
+                new_color = (
+                    random.randint(0, 255),
+                    random.randint(0, 255),
+                    random.randint(0, 255)
+                )
+                if new_color not in colors:
+                    colors.append(new_color)
+            # Use only the number of colors we need
+            selected_colors = random.sample(colors, colors_needed)
+            # Generate variations
+            variations = []
+            for i, color in enumerate(selected_colors):
+                # Apply tint to numpy array
+                tinted_array = self.apply_tint(img_array.copy(), color)
+                variations.append(tinted_array)
+                print(f"Created variation {i+1}/{colors_needed}")
+            print(f"Successfully generated {len(variations)} variations")
+            return variations
+        except Exception as e:
+            print(f"Error in generate_design: {e}")
+            traceback.print_exc()
+            # Return the original image array if there's an error
+            return [np.array(image.convert('RGB'))]

model.py ADDED Viewed

	@@ -0,0 +1,14 @@

+from typing import List
+from PIL import Image
+import numpy as np
+class DesignModel:
+    """Interface for the design model"""
+    def __init__(self):
+        self.seed = None
+        self.neg_prompt = None
+        self.additional_quality_suffix = None
+    def generate_design(self, image: Image.Image, num_variations: int = 1, **kwargs) -> List[np.ndarray]:
+        """Generate design variations from the input image"""
+        raise NotImplementedError("This method should be implemented by concrete model classes")

prod_model.py ADDED Viewed

	@@ -0,0 +1,170 @@

+from model import DesignModel
+from PIL import Image
+import numpy as np
+from typing import List
+import random
+import time
+import torch
+from diffusers import StableDiffusionImg2ImgPipeline
+from transformers import CLIPTokenizer
+import logging
+import os
+from datetime import datetime
+# Set up logging
+log_dir = "logs"
+os.makedirs(log_dir, exist_ok=True)
+log_file = os.path.join(log_dir, f"prod_model_{datetime.now().strftime('%Y%m%d')}.log")
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler(log_file),
+        logging.StreamHandler()
+    ]
+)
+class ProductionDesignModel(DesignModel):
+    def __init__(self):
+        super().__init__()
+        try:
+            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+            logging.info(f"Using device: {self.device}")
+            self.model_id = "stabilityai/stable-diffusion-2-1"
+            logging.info(f"Loading model: {self.model_id}")
+            # Initialize the pipeline with error handling
+            try:
+                self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+                    self.model_id,
+                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                    safety_checker=None  # Disable safety checker for performance
+                ).to(self.device)
+                # Enable optimizations
+                self.pipe.enable_attention_slicing()
+                if self.device == "cuda":
+                    self.pipe.enable_model_cpu_offload()
+                    self.pipe.enable_vae_slicing()
+                logging.info("Model loaded successfully")
+            except Exception as e:
+                logging.error(f"Error loading model: {e}")
+                raise
+            # Initialize tokenizer
+            self.tokenizer = CLIPTokenizer.from_pretrained(self.model_id)
+            # Set default prompts
+            self.neg_prompt = "blurry, low quality, distorted, deformed, disfigured, watermark, text, bad proportions, duplicate, double, multiple, broken, cropped"
+            self.additional_quality_suffix = "interior design, 4K, high resolution, photorealistic"
+        except Exception as e:
+            logging.error(f"Error in initialization: {e}")
+            raise
+    def _prepare_prompt(self, prompt: str) -> str:
+        """Prepare the prompt by adding quality suffix and checking length"""
+        try:
+            full_prompt = f"{prompt}, {self.additional_quality_suffix}"
+            tokens = self.tokenizer.tokenize(full_prompt)
+            if len(tokens) > 77:
+                logging.warning(f"Prompt too long ({len(tokens)} tokens). Truncating...")
+                tokens = tokens[:77]
+                full_prompt = self.tokenizer.convert_tokens_to_string(tokens)
+            logging.info(f"Prepared prompt: {full_prompt}")
+            return full_prompt
+        except Exception as e:
+            logging.error(f"Error preparing prompt: {e}")
+            return prompt  # Return original prompt if processing fails
+    def generate_design(self, image: Image.Image, num_variations: int = 1, **kwargs) -> List[np.ndarray]:
+        """Generate design variations with proper parameter handling"""
+        generation_start = time.time()
+        try:
+            # Log input parameters
+            logging.info(f"Generating {num_variations} variations with parameters: {kwargs}")
+            # Get parameters from kwargs with defaults
+            prompt = kwargs.get('prompt', '')
+            num_steps = int(kwargs.get('num_steps', 50))
+            guidance_scale = float(kwargs.get('guidance_scale', 7.5))
+            strength = float(kwargs.get('strength', 0.75))
+            base_seed = kwargs.get('seed', int(time.time()))
+            # Parameter validation
+            num_steps = max(20, min(100, num_steps))
+            guidance_scale = max(1, min(20, guidance_scale))
+            strength = max(0.1, min(1.0, strength))
+            # Prepare the prompt
+            full_prompt = self._prepare_prompt(prompt)
+            # Generate distinct seeds
+            seeds = [base_seed + i * 10000 for i in range(num_variations)]
+            logging.info(f"Using seeds: {seeds}")
+            # Prepare the input image
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+            # Generate variations
+            variations = []
+            generator = torch.Generator(device=self.device)
+            for i, seed in enumerate(seeds):
+                try:
+                    variation_start = time.time()
+                    generator.manual_seed(seed)
+                    # Generate the image
+                    output = self.pipe(
+                        prompt=full_prompt,
+                        negative_prompt=self.neg_prompt,
+                        image=image,
+                        num_inference_steps=num_steps,
+                        guidance_scale=guidance_scale,
+                        strength=strength,
+                        generator=generator
+                    ).images[0]
+                    variations.append(np.array(output))
+                    variation_time = time.time() - variation_start
+                    logging.info(f"Generated variation {i+1}/{num_variations} in {variation_time:.2f}s")
+                except Exception as e:
+                    logging.error(f"Error generating variation {i+1}: {e}")
+                    if not variations:  # If no successful variations yet
+                        variations.append(np.array(image.convert('RGB')))
+            total_time = time.time() - generation_start
+            logging.info(f"Generation completed in {total_time:.2f}s")
+            return variations
+        except Exception as e:
+            logging.error(f"Error in generate_design: {e}")
+            import traceback
+            logging.error(traceback.format_exc())
+            return [np.array(image.convert('RGB'))]
+        finally:
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+                logging.info("Cleared CUDA cache")
+    def __del__(self):
+        """Cleanup when the model is deleted"""
+        try:
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+                logging.info("Final CUDA cache cleanup")
+        except:
+            pass

requirements.txt CHANGED Viewed

@@ -1,10 +1,33 @@
-diffusers==0.25.0
-xformers==0.0.23.post1
-transformers==4.39.1
-torchvision
-accelerate==0.26.1
-opencv-python==4.9.0.80
-scipy==1.11.4
-triton==2.1.0
-altair==4.1.0
-pandas==2.1.4

+# Core dependencies
+gradio>=3.50.2
+Pillow>=10.0.0
+numpy>=1.24.0
+# Model dependencies
+torch>=2.0.0
+diffusers>=0.21.0
+transformers>=4.31.0
+accelerate>=0.21.0
+# Google Drive integration
+google-auth>=2.22.0
+google-auth-oauthlib>=1.0.0
+google-api-python-client>=2.95.0
+# Utility packages
+python-dateutil>=2.8.2
+tqdm>=4.65.0
+requests>=2.31.0
+# Optional but recommended
+opencv-python>=4.8.0  # For image processing
+safetensors>=0.3.1   # For faster model loading
+# Development tools
+pytest>=7.4.0
+black>=22.0.0
+flake8>=6.0.0
+isort>=5.12.0
+# Testing dependencies
+pytest-mock>=3.11.1

test_prompt.py ADDED Viewed

	@@ -0,0 +1,320 @@

+import unittest
+from app import update_prompt
+from prod_model import ProductionDesignModel
+from PIL import Image
+import numpy as np
+class TestPromptGeneration(unittest.TestCase):
+    def setUp(self):
+        """Set up default values for tests"""
+        self.default_params = {
+            "room": "Living Room",
+            "style": "Modern",
+            "colors": "Neutral",
+            "floor_t": "Keep Existing",
+            "floor_c": "Keep Existing",
+            "floor_p": "Keep Existing",
+            "wall_t": "Keep Existing",
+            "wall_c": "Keep Existing",
+            "wall_f": "Keep Existing",
+            "custom_text": "",
+            "art_en": False,
+            "art_col": "None",
+            "art_size": "None",
+            "mirror_en": False,
+            "mirror_fr": "None",
+            "mirror_size": "Medium",
+            "sconce_en": False,
+            "sconce_col": "None",
+            "sconce_style": "Modern",
+            "shelf_en": False,
+            "shelf_col": "None",
+            "shelf_size": "Medium",
+            "plants_en": False,
+            "plants_type": "None",
+            "plants_size": "Medium"
+        }
+    def test_basic_room_style(self):
+        """Test basic room and style prompt generation"""
+        prompt = update_prompt(**self.default_params)
+        expected = "Design a Modern living room with a Neutral color scheme"
+        self.assertEqual(prompt, expected)
+    def test_all_room_types(self):
+        """Test all room types"""
+        room_types = [
+            "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"
+        ]
+        for room in room_types:
+            params = self.default_params.copy()
+            params["room"] = room
+            prompt = update_prompt(**params)
+            expected = f"Design a Modern {room.lower()} with a Neutral color scheme"
+            self.assertEqual(prompt, expected)
+    def test_all_styles(self):
+        """Test all style presets"""
+        styles = [
+            "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"
+        ]
+        for style in styles:
+            params = self.default_params.copy()
+            params["style"] = style
+            prompt = update_prompt(**params)
+            expected = f"Design a {style} living room with a Neutral color scheme"
+            self.assertEqual(prompt, expected)
+    def test_all_color_schemes(self):
+        """Test all color schemes"""
+        color_schemes = [
+            "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"
+        ]
+        for color in color_schemes:
+            params = self.default_params.copy()
+            params["colors"] = color
+            prompt = update_prompt(**params)
+            expected = f"Design a Modern living room with a {color} color scheme"
+            self.assertEqual(prompt, expected)
+    def test_floor_combinations(self):
+        """Test various floor combinations"""
+        test_cases = [
+            {
+                "floor_t": "Hardwood",
+                "floor_c": "Keep Existing",
+                "floor_p": "Keep Existing",
+                "expected": "featuring Hardwood flooring"
+            },
+            {
+                "floor_t": "Hardwood",
+                "floor_c": "Light Oak",
+                "floor_p": "Keep Existing",
+                "expected": "featuring Hardwood in Light Oak flooring"
+            },
+            {
+                "floor_t": "Hardwood",
+                "floor_c": "Light Oak",
+                "floor_p": "Elegant Herringbone",
+                "expected": "featuring Hardwood in Light Oak with Elegant Herringbone pattern flooring"
+            }
+        ]
+        for case in test_cases:
+            params = self.default_params.copy()
+            params.update({
+                "floor_t": case["floor_t"],
+                "floor_c": case["floor_c"],
+                "floor_p": case["floor_p"]
+            })
+            prompt = update_prompt(**params)
+            expected = f"Design a Modern living room with a Neutral color scheme, {case['expected']}"
+            self.assertEqual(prompt, expected)
+    def test_wall_combinations(self):
+        """Test various wall combinations"""
+        test_cases = [
+            {
+                "wall_t": "Fresh Paint",
+                "wall_c": "Keep Existing",
+                "wall_f": "Keep Existing",
+                "expected": "with Fresh Paint walls"
+            },
+            {
+                "wall_t": "Fresh Paint",
+                "wall_c": "Crisp White",
+                "wall_f": "Keep Existing",
+                "expected": "with Fresh Paint in Crisp White walls"
+            },
+            {
+                "wall_t": "Fresh Paint",
+                "wall_c": "Crisp White",
+                "wall_f": "Pearl Satin",
+                "expected": "with Fresh Paint in Crisp White with Pearl Satin finish walls"
+            }
+        ]
+        for case in test_cases:
+            params = self.default_params.copy()
+            params.update({
+                "wall_t": case["wall_t"],
+                "wall_c": case["wall_c"],
+                "wall_f": case["wall_f"]
+            })
+            prompt = update_prompt(**params)
+            expected = f"Design a Modern living room with a Neutral color scheme, {case['expected']}"
+            self.assertEqual(prompt, expected)
+    def test_accessories_individual(self):
+        """Test each accessory individually"""
+        test_cases = [
+            {
+                "name": "art",
+                "params": {"art_en": True, "art_col": "Vibrant Colors", "art_size": "Oversized"},
+                "expected": "decorated with Oversized Vibrant Colors Art Print"
+            },
+            {
+                "name": "mirror",
+                "params": {"mirror_en": True, "mirror_fr": "Gold", "mirror_size": "Large"},
+                "expected": "decorated with Large Mirror with Gold frame"
+            },
+            {
+                "name": "sconce",
+                "params": {"sconce_en": True, "sconce_col": "Brass", "sconce_style": "Art Deco"},
+                "expected": "decorated with Art Deco Brass Wall Sconce"
+            },
+            {
+                "name": "shelf",
+                "params": {"shelf_en": True, "shelf_col": "Natural Wood", "shelf_size": "Set of 3"},
+                "expected": "decorated with Set of 3 Natural Wood Floating Shelves"
+            },
+            {
+                "name": "plants",
+                "params": {"plants_en": True, "plants_type": "Hanging Plants", "plants_size": "Medium"},
+                "expected": "decorated with Medium Hanging Plants"
+            }
+        ]
+        for case in test_cases:
+            params = self.default_params.copy()
+            params.update(case["params"])
+            prompt = update_prompt(**params)
+            expected = f"Design a Modern living room with a Neutral color scheme, {case['expected']}"
+            self.assertEqual(prompt, expected)
+    def test_custom_text_variations(self):
+        """Test custom text handling"""
+        test_cases = [
+            {"text": "", "should_include": False},
+            {"text": "  ", "should_include": False},
+            {"text": "Add plants", "should_include": True},
+            {"text": "Make it cozy and warm", "should_include": True},
+            {"text": "Multiple\nlines", "should_include": True}
+        ]
+        for case in test_cases:
+            params = self.default_params.copy()
+            params["custom_text"] = case["text"]
+            prompt = update_prompt(**params)
+            base = "Design a Modern living room with a Neutral color scheme"
+            if case["should_include"]:
+                expected = f"{base}, {case['text'].strip()}"
+            else:
+                expected = base
+            self.assertEqual(prompt, expected)
+    def test_complex_combinations(self):
+        """Test complex combinations of all features"""
+        test_cases = [
+            {
+                "name": "full_living_room",
+                "params": {
+                    "room": "Living Room",
+                    "style": "Modern",
+                    "colors": "Warm Gray",
+                    "floor_t": "Hardwood",
+                    "floor_c": "Light Oak",
+                    "floor_p": "Elegant Herringbone",
+                    "wall_t": "Fresh Paint",
+                    "wall_c": "Crisp White",
+                    "wall_f": "Pearl Satin",
+                    "custom_text": "Make it perfect for entertaining",
+                    "art_en": True,
+                    "art_col": "Modern Abstract",
+                    "art_size": "Statement",
+                    "mirror_en": True,
+                    "mirror_fr": "Gold",
+                    "mirror_size": "Large",
+                    "sconce_en": True,
+                    "sconce_col": "Brass",
+                    "sconce_style": "Art Deco",
+                    "shelf_en": True,
+                    "shelf_col": "Natural Wood",
+                    "shelf_size": "Set of 3",
+                    "plants_en": True,
+                    "plants_type": "Hanging Plants",
+                    "plants_size": "Medium"
+                }
+            },
+            {
+                "name": "minimal_bedroom",
+                "params": {
+                    "room": "Bedroom",
+                    "style": "Japanese",
+                    "colors": "Minimalist White",
+                    "floor_t": "Natural Bamboo",
+                    "floor_c": "Keep Existing",
+                    "floor_p": "Keep Existing",
+                    "wall_t": "Fresh Paint",
+                    "wall_c": "Soft White",
+                    "wall_f": "Keep Existing",
+                    "custom_text": "Focus on minimalism and zen aesthetics"
+                }
+            }
+        ]
+        for case in test_cases:
+            params = self.default_params.copy()
+            params.update(case["params"])
+            prompt = update_prompt(**params)
+            self.assertTrue(len(prompt) > 0)
+            self.assertTrue(prompt.startswith("Design a"))
+class TestProductionModel(unittest.TestCase):
+    def setUp(self):
+        """Set up test environment"""
+        self.model = ProductionDesignModel()
+        # Create a simple test image
+        self.test_image = Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8))
+    def test_number_of_variations(self):
+        """Test that the model correctly handles different numbers of variations"""
+        test_cases = [1, 3, 10, 25, 50]  # Test various numbers of variations
+        for num_variations in test_cases:
+            variations = self.model.generate_design(
+                image=self.test_image,
+                num_variations=num_variations,
+                prompt="Test prompt",
+                num_steps=20,  # Minimum steps for faster testing
+                guidance_scale=7.5,
+                strength=0.75
+            )
+            self.assertEqual(
+                len(variations),
+                num_variations,
+                f"Expected {num_variations} variations, got {len(variations)}"
+            )
+    def test_invalid_variation_numbers(self):
+        """Test handling of invalid numbers of variations"""
+        test_cases = [-1, 0, 51, 100]  # Test invalid numbers
+        for num_variations in test_cases:
+            variations = self.model.generate_design(
+                image=self.test_image,
+                num_variations=num_variations,
+                prompt="Test prompt",
+                num_steps=20,
+                guidance_scale=7.5,
+                strength=0.75
+            )
+            # Should clamp to valid range (1-50)
+            self.assertTrue(
+                1 <= len(variations) <= 50,
+                f"Number of variations {len(variations)} outside valid range 1-50"
+            )
+if __name__ == '__main__':
+    unittest.main()