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