Enhance UI dropdown options and improve ProductionDesignModel initialization

- Updated UI dropdowns in app.py to provide a comprehensive list of choices for room types, design styles, and color moods, enhancing user experience.
- Refactored layout for better organization of UI elements.
- Improved the ProductionDesignModel class in prod_model.py by implementing a more robust model initialization process, including advanced architecture setup and detailed logging for better error tracking.
- Added new model dependencies in requirements.txt to support the updated functionality.

app.py

@@ -256,29 +256,53 @@ def create_ui(model: DesignModel):
                 with gr.Group():
                     gr.Markdown("## 🏠 Basic Settings")
                     with gr.Row():
-                        room_type = gr.Dropdown(
+                    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",
-                            choices=["None"] + ["Living Room", "Bedroom", "Dining Room", "Kitchen", "Bathroom", "Home Office", "Master Bedroom", "Guest Room", "Study Room", "Game Room", "Media Room", "Nursery", "Gym", "Library"],
-                            value="None"
-                        )
-                        style_preset = gr.Dropdown(
+                        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",
-                            choices=["None"] + ["Modern", "Contemporary", "Minimalist", "Industrial", "Scandinavian", "Mid-Century Modern", "Traditional", "Transitional", "Farmhouse", "Rustic", "Bohemian", "Art Deco", "Coastal", "Mediterranean", "Japanese", "French Country", "Victorian", "Colonial", "Gothic", "Baroque", "Rococo", "Neoclassical", "Eclectic", "Zen", "Tropical", "Shabby Chic", "Hollywood Regency", "Southwestern", "Asian Fusion", "Retro"],
-                            value="None"
-                        )
-                        color_scheme = gr.Dropdown(
+                            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",
-                            choices=["None"] + ["Neutral", "Monochromatic", "Minimalist White", "Warm Gray", "Cool Gray", "Earth Tones", "Pastel", "Bold Primary", "Jewel Tones", "Black and White", "Navy and Gold", "Forest Green", "Desert Sand", "Ocean Blue", "Sunset Orange", "Deep Purple", "Emerald Green", "Ruby Red", "Sapphire Blue", "Golden Yellow", "Sage Green", "Dusty Rose", "Charcoal", "Cream", "Burgundy", "Teal", "Copper", "Silver", "Bronze", "Slate"],
-                            value="None"
+                            value="Neutral"
                         )
 
             # Row 2 - Surface Finishes
-            with gr.Row():
+                        with gr.Row():
                 # Floor Options
                 with gr.Column(scale=1):
                     with gr.Group():
                         gr.Markdown("## 🎨 Floor Options")
-                        floor_type = gr.Dropdown(
+                            floor_type = gr.Dropdown(
                             choices=[
                                 "Keep Existing", "Hardwood", "Stone Tiles", "Porcelain Tiles",
                                 "Soft Carpet", "Polished Concrete", "Marble", "Vinyl",
@@ -287,9 +311,9 @@ def create_ui(model: DesignModel):
                                 "Mosaic Tiles", "Luxury Vinyl Tiles", "Stained Concrete"
                             ],
                             label="Material",
-                            value="Keep Existing"
-                        )
-                        floor_color = gr.Dropdown(
+                                value="Keep Existing"
+                            )
+                            floor_color = gr.Dropdown(
                             choices=[
                                 "Keep Existing", "Light Oak", "Rich Walnut", "Cool Gray",
                                 "Whitewashed", "Warm Cherry", "Deep Brown", "Classic Black",
@@ -298,10 +322,10 @@ def create_ui(model: DesignModel):
                                 "Cream Travertine", "Dark Slate", "Golden Teak",
                                 "Rustic Pine", "Ebony"
                             ],
-                            label="Color",
-                            value="Keep Existing"
-                        )
-                        floor_pattern = gr.Dropdown(
+                                label="Color",
+                                value="Keep Existing"
+                            )
+                            floor_pattern = gr.Dropdown(
                             choices=[
                                 "Keep Existing", "Classic Straight", "Elegant Herringbone",
                                 "V-Pattern", "Decorative Parquet", "Diagonal Layout",
@@ -311,17 +335,17 @@ def create_ui(model: DesignModel):
                                 "Windmill Pattern", "Large Format", "Mixed Width"
                             ],
                             label="Pattern",
-                            value="Keep Existing"
-                        )
-                
+                                value="Keep Existing"
+                            )
+
                 # Wall Options
                 with gr.Column(scale=1):
                     with gr.Group():
                         gr.Markdown("## 🎨 Wall Options")
-                        wall_type = gr.Dropdown(
+                            wall_type = gr.Dropdown(
                             choices=[
                                 "Keep Existing", "Fresh Paint", "Designer Wallpaper",
-                                "Textured Finish", "Wood Panels", "Exposed Brick",
+                                       "Textured Finish", "Wood Panels", "Exposed Brick",
                                 "Natural Stone", "Wooden Planks", "Modern Concrete",
                                 "Venetian Plaster", "Wainscoting", "Shiplap",
                                 "3D Wall Panels", "Fabric Panels", "Metal Panels",
@@ -329,9 +353,9 @@ def create_ui(model: DesignModel):
                                 "Acoustic Panels", "Living Wall"
                             ],
                             label="Treatment",
-                            value="Keep Existing"
-                        )
-                        wall_color = gr.Dropdown(
+                                value="Keep Existing"
+                            )
+                            wall_color = gr.Dropdown(
                             choices=[
                                 "Keep Existing", "Crisp White", "Soft White", "Warm Beige",
                                 "Gentle Gray", "Sky Blue", "Nature Green", "Sunny Yellow",
@@ -339,10 +363,10 @@ def create_ui(model: DesignModel):
                                 "Terracotta", "Navy Blue", "Charcoal Gray", "Lavender",
                                 "Olive Green", "Dusty Rose", "Teal", "Burgundy"
                             ],
-                            label="Color",
-                            value="Keep Existing"
-                        )
-                        wall_finish = gr.Dropdown(
+                                label="Color",
+                                value="Keep Existing"
+                            )
+                            wall_finish = gr.Dropdown(
                             choices=[
                                 "Keep Existing", "Soft Matte", "Subtle Eggshell",
                                 "Pearl Satin", "Sleek Semi-Gloss", "High Gloss",
@@ -351,9 +375,9 @@ def create_ui(model: DesignModel):
                                 "Venetian", "Lime Wash", "Concrete", "Rustic",
                                 "Lacquered", "Hammered", "Patina"
                             ],
-                            label="Finish",
-                            value="Keep Existing"
-                        )
+                                label="Finish",
+                                value="Keep Existing"
+                            )
 
             # Row 3 - Wall Decorations and Special Requests
             with gr.Row(elem_classes="wall-decorations-row"):
@@ -367,7 +391,7 @@ def create_ui(model: DesignModel):
                             with gr.Column():
                                 with gr.Row():
                                     art_print_enable = gr.Checkbox(label="Add Artwork", value=False)
-                                art_print_color = gr.Dropdown(
+                            art_print_color = gr.Dropdown(
                                     choices=[
                                         "None", "Classic Black & White", "Vibrant Colors",
                                         "Single Color", "Soft Colors", "Modern Abstract",
@@ -378,8 +402,8 @@ def create_ui(model: DesignModel):
                                     ],
                                     label="Art Style",
                                     value="None"
-                                )
-                                art_print_size = gr.Dropdown(
+                            )
+                            art_print_size = gr.Dropdown(
                                     choices=[
                                         "None", "Modest", "Standard", "Statement", "Oversized",
                                         "Gallery Wall", "Diptych", "Triptych", "Mini Series",
@@ -391,9 +415,9 @@ def create_ui(model: DesignModel):
                             
                             # Mirror
                             with gr.Column():
-                                with gr.Row():
+                        with gr.Row():
                                     mirror_enable = gr.Checkbox(label="Add Mirror", value=False)
-                                mirror_frame = gr.Dropdown(
+                            mirror_frame = gr.Dropdown(
                                     choices=[
                                         "None", "Gold", "Silver", "Black", "White", "Wood",
                                         "Brass", "Bronze", "Copper", "Chrome", "Antique Gold",
@@ -402,8 +426,8 @@ def create_ui(model: DesignModel):
                                     ],
                                     label="Frame Style",
                                     value="None"
-                                )
-                                mirror_size = gr.Dropdown(
+                            )
+                            mirror_size = gr.Dropdown(
                                     choices=[
                                         "Small", "Medium", "Large", "Full Length",
                                         "Oversized", "Double Width", "Floor Mirror",
@@ -419,7 +443,7 @@ def create_ui(model: DesignModel):
                             with gr.Column():
                                 with gr.Row():
                                     sconce_enable = gr.Checkbox(label="Add Wall Sconce", value=False)
-                                sconce_color = gr.Dropdown(
+                            sconce_color = gr.Dropdown(
                                     choices=[
                                         "None", "Black", "Gold", "Silver", "Bronze", "White",
                                         "Brass", "Copper", "Chrome", "Antique Brass",
@@ -429,8 +453,8 @@ def create_ui(model: DesignModel):
                                     ],
                                     label="Sconce Color",
                                     value="None"
-                                )
-                                sconce_style = gr.Dropdown(
+                            )
+                            sconce_style = gr.Dropdown(
                                     choices=[
                                         "Modern", "Traditional", "Industrial", "Art Deco",
                                         "Minimalist", "Vintage", "Contemporary", "Rustic",
@@ -444,9 +468,9 @@ def create_ui(model: DesignModel):
 
                             # Floating Shelves
                             with gr.Column():
-                                with gr.Row():
+                        with gr.Row():
                                     shelf_enable = gr.Checkbox(label="Add Floating Shelves", value=False)
-                                shelf_color = gr.Dropdown(
+                            shelf_color = gr.Dropdown(
                                     choices=[
                                         "None", "White", "Black", "Natural Wood", "Glass",
                                         "Dark Wood", "Light Wood", "Metal", "Gold", "Silver",
@@ -456,8 +480,8 @@ def create_ui(model: DesignModel):
                                     ],
                                     label="Shelf Material",
                                     value="None"
-                                )
-                                shelf_size = gr.Dropdown(
+                            )
+                            shelf_size = gr.Dropdown(
                                     choices=[
                                         "Small", "Medium", "Large", "Set of 3",
                                         "Extra Long", "Corner Set", "Asymmetric Set",
@@ -466,13 +490,13 @@ def create_ui(model: DesignModel):
                                     ],
                                     label="Shelf Size",
                                     value="Medium"
-                                )
+                            )
 
-                            # Plants
+                        # Plants
                             with gr.Column():
-                                with gr.Row():
+                        with gr.Row():
                                     plants_enable = gr.Checkbox(label="Add Plants", value=False)
-                                plants_type = gr.Dropdown(
+                            plants_type = gr.Dropdown(
                                     choices=[
                                         "None", "Hanging Plants", "Vertical Garden",
                                         "Plant Shelf", "Single Plant", "Climbing Vines",
@@ -483,8 +507,8 @@ def create_ui(model: DesignModel):
                                     ],
                                     label="Plant Type",
                                     value="None"
-                                )
-                                plants_size = gr.Dropdown(
+                            )
+                            plants_size = gr.Dropdown(
                                     choices=[
                                         "Small", "Medium", "Large", "Mixed Sizes",
                                         "Full Wall", "Statement Piece", "Compact",
@@ -499,7 +523,7 @@ def create_ui(model: DesignModel):
                 with gr.Column(scale=1):
                     with gr.Group():
                         gr.Markdown("## ✨ Special Requests")
-                        input_text = gr.Textbox(
+                    input_text = gr.Textbox(
                             label="Additional Details",
                             placeholder="Add any special requests or details here...",
                             lines=3
@@ -517,14 +541,14 @@ def create_ui(model: DesignModel):
                             step=1,
                             label="Quality Steps"
                         )
-                        guidance_scale = gr.Slider(
+                            guidance_scale = gr.Slider(
                             minimum=1,
                             maximum=20,
                             value=7.5,
                             step=0.1,
                             label="Design Freedom"
-                        )
-                        strength = gr.Slider(
+                            )
+                            strength = gr.Slider(
                             minimum=0.1,
                             maximum=1.0,
                             value=0.75,
@@ -544,7 +568,7 @@ def create_ui(model: DesignModel):
                             )
 
             # Row 4 - Current Prompts
-            with gr.Row():
+                    with gr.Row():
                 with gr.Group():
                     gr.Markdown("## 📝 Current Prompts")
                     prompt_display = gr.TextArea(
@@ -858,9 +882,9 @@ def main():
     is_test_mode = "--test" in sys.argv
     
     if is_test_mode:
-        print("Starting in TEST mode...")
+    print("Starting in TEST mode...")
         from mock_model import MockDesignModel
-        model = MockDesignModel()
+    model = MockDesignModel()
     else:
         print("Starting in PRODUCTION mode...")
         from prod_model import ProductionDesignModel

prod_model.py

@@ -5,11 +5,13 @@ from typing import List
 import random
 import time
 import torch
-from diffusers import StableDiffusionImg2ImgPipeline
-from transformers import CLIPTokenizer
+from diffusers.pipelines.controlnet import StableDiffusionControlNetInpaintPipeline
+from diffusers import ControlNetModel, UniPCMultistepScheduler, AutoPipelineForText2Image
+from transformers import AutoImageProcessor, UperNetForSemanticSegmentation, AutoModelForDepthEstimation
 import logging
 import os
 from datetime import datetime
+import gc
 
 # Set up logging
 log_dir = "logs"
@@ -27,158 +29,217 @@ logging.basicConfig(
 
 class ProductionDesignModel(DesignModel):
     def __init__(self):
-        super().__init__()
+        """Initialize the production model with advanced architecture"""
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+        
+        # Setup logging
+        logging.basicConfig(filename=f'logs/prod_model_{time.strftime("%Y%m%d")}.log',
+                          level=logging.INFO,
+                          format='%(asctime)s - %(levelname)s - %(message)s')
+        
+        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"
+        
         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"
-            self.tokenizer_id = "openai/clip-vit-large-patch14"  # Correct tokenizer for SD 2.1
-            logging.info(f"Loading model: {self.model_id}")
-            logging.info(f"Loading tokenizer: {self.tokenizer_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 with correct path
-            try:
-                self.tokenizer = CLIPTokenizer.from_pretrained(self.tokenizer_id)
-                logging.info("Tokenizer loaded successfully")
-            except Exception as e:
-                logging.error(f"Error loading tokenizer: {e}")
-                raise
-            
-            # 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"
-            
+            logging.info(f"Initializing models on {self.device} with {self.dtype}")
+            self._initialize_models()
+            logging.info("Models initialized successfully")
         except Exception as e:
-            logging.error(f"Error in initialization: {e}")
+            logging.error(f"Error initializing models: {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
+    def _initialize_models(self):
+        """Initialize all required models and pipelines"""
+        # Initialize ControlNet models
+        self.controlnet_depth = ControlNetModel.from_pretrained(
+            "controlnet_depth", torch_dtype=self.dtype, use_safetensors=True
+        )
+        self.controlnet_seg = ControlNetModel.from_pretrained(
+            "own_controlnet", torch_dtype=self.dtype, use_safetensors=True
+        )
+
+        # Initialize main pipeline
+        self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+            "SG161222/Realistic_Vision_V5.1_noVAE",
+            controlnet=[self.controlnet_depth, self.controlnet_seg],
+            safety_checker=None,
+            torch_dtype=self.dtype
+        )
+
+        # Setup IP-Adapter
+        self.pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models",
+                                weight_name="ip-adapter_sd15.bin")
+        self.pipe.set_ip_adapter_scale(0.4)
+        self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
+        self.pipe = self.pipe.to(self.device)
+
+        # Initialize guide pipeline
+        self.guide_pipe = AutoPipelineForText2Image.from_pretrained(
+            "segmind/SSD-1B",
+            torch_dtype=self.dtype,
+            use_safetensors=True,
+            variant="fp16"
+        ).to(self.device)
+
+        # Initialize segmentation and depth models
+        self.seg_processor, self.seg_model = self._init_segmentation()
+        self.depth_processor, self.depth_model = self._init_depth()
+        self.depth_model = self.depth_model.to(self.device)
+
+    def _init_segmentation(self):
+        """Initialize segmentation models"""
+        processor = AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-small")
+        model = UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-small")
+        return processor, model
+
+    def _init_depth(self):
+        """Initialize depth estimation models"""
+        processor = AutoImageProcessor.from_pretrained(
+            "LiheYoung/depth-anything-large-hf",
+            torch_dtype=self.dtype
+        )
+        model = AutoModelForDepthEstimation.from_pretrained(
+            "LiheYoung/depth-anything-large-hf",
+            torch_dtype=self.dtype
+        )
+        return processor, model
+
+    def _get_depth_map(self, image: Image.Image) -> Image.Image:
+        """Generate depth map for input image"""
+        image_to_depth = self.depth_processor(images=image, return_tensors="pt").to(self.device)
+        with torch.inference_mode():
+            depth_map = self.depth_model(**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]
+        return Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
+
+    def _segment_image(self, image: Image.Image) -> Image.Image:
+        """Generate segmentation map for input image"""
+        pixel_values = self.seg_processor(image, return_tensors="pt").pixel_values
+        with torch.inference_mode():
+            outputs = self.seg_model(pixel_values)
+
+        seg = self.seg_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)
+        
+        # You'll need to implement the palette mapping here
+        # This is a placeholder - you should implement proper color mapping
+        for label in range(seg.max() + 1):
+            color_seg[seg == label, :] = [label * 30 % 255] * 3
             
-        except Exception as e:
-            logging.error(f"Error preparing prompt: {e}")
-            return prompt  # Return original prompt if processing fails
+        return Image.fromarray(color_seg).convert('RGB')
 
-    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()
+    def _resize_image(self, image: Image.Image, target_size: int) -> Image.Image:
+        """Resize image while maintaining aspect ratio"""
+        width, height = image.size
+        if width > height:
+            new_width = target_size
+            new_height = int(height * (target_size / width))
+        else:
+            new_height = target_size
+            new_width = int(width * (target_size / height))
+        return image.resize((new_width, new_height), Image.LANCZOS)
+
+    def _flush(self):
+        """Clear CUDA cache"""
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def generate_design(self, image: Image.Image, prompt: str, **kwargs) -> List[Image.Image]:
+        """
+        Generate design variations based on input image and prompt
+        """
         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))
-            
-            # Handle seed properly
+            # Set seed
             seed_param = kwargs.get('seed')
             base_seed = int(time.time()) if seed_param is None else int(seed_param)
-            logging.info(f"Using base seed: {base_seed}")
-            
-            # 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))
-            
-            # Log validated parameters
-            logging.info(f"Validated parameters: steps={num_steps}, guidance={guidance_scale}, strength={strength}")
+            self.generator = torch.Generator(device=self.device).manual_seed(base_seed)
             
-            # 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}")
+            # Get parameters
+            num_variations = kwargs.get('num_variations', 1)
+            guidance_scale = float(kwargs.get('guidance_scale', 10.0))
+            num_steps = int(kwargs.get('num_steps', 50))
+            strength = float(kwargs.get('strength', 0.9))
+            img_size = int(kwargs.get('img_size', 768))
+
+            logging.info(f"Generating design with parameters: guidance_scale={guidance_scale}, "
+                        f"num_steps={num_steps}, strength={strength}, img_size={img_size}")
+
+            # Prepare prompt
+            pos_prompt = f"{prompt}, {self.additional_quality_suffix}"
+
+            # Process input image
+            orig_size = image.size
+            input_image = self._resize_image(image, img_size)
             
-            # Prepare the input image
-            if image.mode != "RGB":
-                image = image.convert("RGB")
+            # Generate depth map
+            depth_map = self._get_depth_map(input_image)
             
+            # Generate segmentation
+            seg_map = self._segment_image(input_image)
+
+            # Generate IP-adapter reference image
+            self._flush()
+            ip_image = self.guide_pipe(
+                pos_prompt,
+                num_inference_steps=num_steps,
+                negative_prompt=self.neg_prompt,
+                generator=self.generator
+            ).images[0]
+
             # Generate variations
             variations = []
-            generator = torch.Generator(device=self.device)
-            
-            for i, seed in enumerate(seeds):
+            for i in range(num_variations):
                 try:
-                    variation_start = time.time()
-                    generator.manual_seed(seed)
-                    
-                    # Generate the image
-                    output = self.pipe(
-                        prompt=full_prompt,
+                    self._flush()
+                    variation = self.pipe(
+                        prompt=pos_prompt,
                         negative_prompt=self.neg_prompt,
-                        image=image,
                         num_inference_steps=num_steps,
-                        guidance_scale=guidance_scale,
                         strength=strength,
-                        generator=generator
+                        guidance_scale=guidance_scale,
+                        generator=self.generator,
+                        image=input_image,
+                        ip_adapter_image=ip_image,
+                        control_image=[depth_map, seg_map],
+                        controlnet_conditioning_scale=[0.5, 0.5]
                     ).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")
+                    # Resize back to original size
+                    variation = variation.resize(orig_size, Image.LANCZOS)
+                    variations.append(variation)
                     
                 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")
-            
+                    logging.error(f"Error generating variation {i}: {e}")
+                    continue
+
+            if not variations:
+                logging.warning("No variations were generated successfully")
+                return [image]  # Return original image if no variations were generated
+
             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")
-        
+            return [image]  # Return original image in case of error
+
     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 
+        self._flush() 

requirements.txt

@@ -2,32 +2,24 @@
 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
+safetensors>=0.3.1
 
 # Google Drive integration
-google-auth>=2.22.0
-google-auth-oauthlib>=1.0.0
 google-api-python-client>=2.95.0
+google-auth-oauthlib>=1.0.0
+google-auth>=2.22.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
+tqdm>=4.65.0
+opencv-python>=4.8.0
 
 # 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
+pytest-mock>=3.11.1
+mock>=5.1.0