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'))]