Spaces:
Sleeping
Sleeping
File size: 3,257 Bytes
4e4b650 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 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'))] |