Spaces:
Running
Running
| 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'))] |