Initial commit: Add project structure and core modules

- Add color palette extraction with K-means clustering
- Add color theory harmonies (complementary, triadic, analogous, etc.)
- Add palette swapping with nearest neighbor matching
- Add palette visualization tools
- Set up project structure with README and requirements

.gitignore

@@ -0,0 +1,51 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+dist/
+*.egg-info/
+.eggs/
+
+# Virtual Environment
+venv/
+env/
+ENV/
+.venv
+
+# Environment variables
+.env
+.env.local
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# OS
+.DS_Store
+.DS_Store?
+._*
+Thumbs.db
+
+# Gradio
+gradio_cached_examples/
+flagged/
+
+# Generated images
+generated_images/
+output/
+temp/
+
+# Model cache
+transformers_cache/
+huggingface_cache/
+diffusers_cache/
+
+# Logs
+*.log
+logs/

README.md

@@ -0,0 +1,35 @@
+# AI Image Editor
+
+An AI-powered image editor with color palette extraction and swapping capabilities.
+
+## Features
+
+- **AI Image Generation**: Create images from text prompts using Stable Diffusion
+- **Color Palette Extraction**: Automatically extract dominant colors from images
+- **Palette Swapping**: Replace colors in images with custom palettes
+- **Color Theory Tools**: Apply color harmony rules (complementary, analogous, triadic)
+- **Interactive UI**: Easy-to-use Gradio interface
+
+## Installation
+
+```bash
+pip install -r requirements.txt
+```
+
+## Usage
+
+```bash
+python app.py
+```
+
+## Tech Stack
+
+- **Gradio**: Web interface
+- **Diffusers**: Stable Diffusion image generation
+- **Pillow**: Image processing
+- **scikit-learn**: K-means clustering for palette extraction
+- **NumPy**: Numerical operations
+
+## License
+
+MIT

color_palette.py

@@ -0,0 +1,329 @@
+"""
+Color Palette Extraction and Manipulation Module
+Implements patterns from Pylette, colorgram, and paletteswapper
+"""
+
+import numpy as np
+from PIL import Image
+from sklearn.cluster import KMeans
+from typing import List, Tuple
+import colorsys
+
+
+class ColorPalette:
+    """
+    Handle color palette extraction and manipulation
+    Based on research from top GitHub repos:
+    - Pylette: KMeans extraction with metadata
+    - colorgram.py: Fast extraction with proportions
+    - paletteswapper: Clean functional palette replacement
+    """
+
+    @staticmethod
+    def extract_palette(
+        image: Image.Image,
+        n_colors: int = 5,
+        sort_by: str = 'frequency'  # 'frequency' or 'luminance'
+    ) -> List[Tuple[int, int, int]]:
+        """
+        Extract dominant colors using K-means clustering
+        Pattern from Pylette and scikit-learn tutorials
+
+        Args:
+            image: PIL Image
+            n_colors: Number of colors to extract
+            sort_by: How to sort results ('frequency' or 'luminance')
+
+        Returns:
+            List of RGB tuples sorted by specified criterion
+        """
+        # Resize for faster processing (pattern from colorgram.py)
+        img = image.copy()
+        img.thumbnail((200, 200))
+        img = img.convert('RGB')
+
+        # Flatten image to pixel array
+        pixels = np.array(img).reshape(-1, 3)
+
+        # K-means clustering (industry standard from scikit-learn)
+        kmeans = KMeans(
+            n_clusters=n_colors,
+            random_state=42,
+            n_init=10,
+            max_iter=300
+        )
+        kmeans.fit(pixels)
+
+        # Get cluster centers (dominant colors)
+        colors = kmeans.cluster_centers_.astype(int)
+        labels = kmeans.labels_
+        counts = np.bincount(labels)
+
+        # Calculate proportions (pattern from colorgram.py)
+        total_pixels = len(labels)
+        proportions = counts / total_pixels
+
+        # Sort based on criterion
+        if sort_by == 'luminance':
+            # Calculate luminance using standard formula
+            luminances = 0.299 * colors[:, 0] + 0.587 * colors[:, 1] + 0.114 * colors[:, 2]
+            sorted_indices = np.argsort(-luminances)
+        else:  # frequency
+            sorted_indices = np.argsort(-counts)
+
+        sorted_colors = colors[sorted_indices]
+
+        return [tuple(color) for color in sorted_colors]
+
+    @staticmethod
+    def swap_palette(
+        image: Image.Image,
+        source_palette: List[Tuple[int, int, int]],
+        target_palette: List[Tuple[int, int, int]],
+        mode: str = 'closest'  # 'closest' or 'threshold'
+    ) -> Image.Image:
+        """
+        Replace colors in image from source to target palette
+        Pattern from paletteswapper with nearest neighbor matching
+
+        Args:
+            image: PIL Image to modify
+            source_palette: Original colors to replace
+            target_palette: New colors to use
+            mode: 'closest' for all pixels, 'threshold' for selective
+
+        Returns:
+            Modified PIL Image
+        """
+        img_array = np.array(image.convert('RGB')).astype(float)
+        height, width, _ = img_array.shape
+        pixels = img_array.reshape(-1, 3)
+        result_pixels = pixels.copy()
+
+        # Map each source color to target color
+        min_len = min(len(source_palette), len(target_palette))
+
+        for i in range(min_len):
+            source_color = np.array(source_palette[i])
+            target_color = np.array(target_palette[i])
+
+            # Calculate Euclidean distance to source color
+            distances = np.linalg.norm(pixels - source_color, axis=1)
+
+            if mode == 'threshold':
+                # Only replace pixels very close to source color
+                threshold = np.percentile(distances, 15)
+                mask = distances <= threshold
+            else:  # closest
+                # For each pixel, check if this is the closest source color
+                all_distances = np.array([
+                    np.linalg.norm(pixels - np.array(src), axis=1)
+                    for src in source_palette[:min_len]
+                ])
+                closest_color_idx = np.argmin(all_distances, axis=0)
+                mask = closest_color_idx == i
+
+            result_pixels[mask] = target_color
+
+        # Reshape and convert back to image
+        result_array = result_pixels.reshape(height, width, 3).astype(np.uint8)
+        return Image.fromarray(result_array)
+
+    @staticmethod
+    def rgb_to_hex(rgb: Tuple[int, int, int]) -> str:
+        """Convert RGB tuple to hex string"""
+        return '#{:02x}{:02x}{:02x}'.format(*rgb)
+
+    @staticmethod
+    def hex_to_rgb(hex_color: str) -> Tuple[int, int, int]:
+        """Convert hex string to RGB tuple"""
+        hex_color = hex_color.lstrip('#')
+        return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
+
+
+class ColorTheory:
+    """
+    Color harmony generation using color theory
+    Implements patterns from colorharmonies library
+    """
+
+    @staticmethod
+    def complementary(base_color: Tuple[int, int, int]) -> List[Tuple[int, int, int]]:
+        """
+        Generate complementary palette (opposite on color wheel)
+        180° rotation in HSV space
+        """
+        r, g, b = [c / 255.0 for c in base_color]
+        h, s, v = colorsys.rgb_to_hsv(r, g, b)
+
+        colors = [base_color]
+
+        # Complementary: rotate hue by 180°
+        comp_h = (h + 0.5) % 1.0
+        comp_rgb = colorsys.hsv_to_rgb(comp_h, s, v)
+        colors.append(tuple(int(c * 255) for c in comp_rgb))
+
+        return colors
+
+    @staticmethod
+    def analogous(base_color: Tuple[int, int, int]) -> List[Tuple[int, int, int]]:
+        """
+        Generate analogous palette (adjacent on color wheel)
+        ±30° rotation in HSV space
+        """
+        r, g, b = [c / 255.0 for c in base_color]
+        h, s, v = colorsys.rgb_to_hsv(r, g, b)
+
+        colors = []
+
+        # Three analogous colors: -30°, base, +30°
+        for offset in [-1/12, 0, 1/12]:  # ±30° = ±1/12 of circle
+            new_h = (h + offset) % 1.0
+            new_rgb = colorsys.hsv_to_rgb(new_h, s, v)
+            colors.append(tuple(int(c * 255) for c in new_rgb))
+
+        return colors
+
+    @staticmethod
+    def triadic(base_color: Tuple[int, int, int]) -> List[Tuple[int, int, int]]:
+        """
+        Generate triadic palette (evenly spaced on color wheel)
+        120° spacing in HSV space
+        """
+        r, g, b = [c / 255.0 for c in base_color]
+        h, s, v = colorsys.rgb_to_hsv(r, g, b)
+
+        colors = []
+
+        # Three colors at 0°, 120°, 240°
+        for offset in [0, 1/3, 2/3]:
+            new_h = (h + offset) % 1.0
+            new_rgb = colorsys.hsv_to_rgb(new_h, s, v)
+            colors.append(tuple(int(c * 255) for c in new_rgb))
+
+        return colors
+
+    @staticmethod
+    def split_complementary(base_color: Tuple[int, int, int]) -> List[Tuple[int, int, int]]:
+        """
+        Generate split-complementary palette
+        Base + two colors adjacent to complement (±150°)
+        """
+        r, g, b = [c / 255.0 for c in base_color]
+        h, s, v = colorsys.rgb_to_hsv(r, g, b)
+
+        colors = [base_color]
+
+        # Two colors at ±150° from base (30° on each side of complement)
+        for offset in [5/12, 7/12]:  # 150° and 210° = 5/12 and 7/12 of circle
+            new_h = (h + offset) % 1.0
+            new_rgb = colorsys.hsv_to_rgb(new_h, s, v)
+            colors.append(tuple(int(c * 255) for c in new_rgb))
+
+        return colors
+
+    @staticmethod
+    def tetradic(base_color: Tuple[int, int, int]) -> List[Tuple[int, int, int]]:
+        """
+        Generate tetradic/double-complementary palette
+        Two complementary pairs (90° spacing)
+        """
+        r, g, b = [c / 255.0 for c in base_color]
+        h, s, v = colorsys.rgb_to_hsv(r, g, b)
+
+        colors = []
+
+        # Four colors at 0°, 90°, 180°, 270°
+        for offset in [0, 0.25, 0.5, 0.75]:
+            new_h = (h + offset) % 1.0
+            new_rgb = colorsys.hsv_to_rgb(new_h, s, v)
+            colors.append(tuple(int(c * 255) for c in new_rgb))
+
+        return colors
+
+    @staticmethod
+    def monochromatic(base_color: Tuple[int, int, int], n_colors: int = 5) -> List[Tuple[int, int, int]]:
+        """
+        Generate monochromatic palette (variations of one hue)
+        Varies saturation and value while keeping hue constant
+        """
+        r, g, b = [c / 255.0 for c in base_color]
+        h, s, v = colorsys.rgb_to_hsv(r, g, b)
+
+        colors = []
+
+        for i in range(n_colors):
+            # Vary value (brightness) from dark to light
+            new_v = 0.2 + (i / (n_colors - 1)) * 0.7  # 20% to 90%
+            # Slightly vary saturation for visual interest
+            new_s = max(0.3, min(1.0, s + (i / n_colors - 0.5) * 0.3))
+            new_rgb = colorsys.hsv_to_rgb(h, new_s, new_v)
+            colors.append(tuple(int(c * 255) for c in new_rgb))
+
+        return colors
+
+
+class PaletteVisualizer:
+    """Create visual representations of color palettes"""
+
+    @staticmethod
+    def create_palette_image(
+        palette: List[Tuple[int, int, int]],
+        width: int = 500,
+        height: int = 100,
+        show_hex: bool = False
+    ) -> Image.Image:
+        """
+        Create a visual swatch of colors
+        Pattern from Pylette visualization
+        """
+        n_colors = len(palette)
+        if n_colors == 0:
+            return Image.new('RGB', (width, height), (255, 255, 255))
+
+        color_width = width // n_colors
+
+        palette_img = Image.new('RGB', (width, height))
+        pixels = palette_img.load()
+
+        for i, color in enumerate(palette):
+            x_start = i * color_width
+            x_end = (i + 1) * color_width if i < n_colors - 1 else width
+
+            for x in range(x_start, x_end):
+                for y in range(height):
+                    pixels[x, y] = color
+
+        return palette_img
+
+    @staticmethod
+    def create_comparison_image(
+        original: Image.Image,
+        modified: Image.Image,
+        original_palette: List[Tuple[int, int, int]],
+        target_palette: List[Tuple[int, int, int]]
+    ) -> Image.Image:
+        """Create side-by-side comparison with palettes"""
+        # Create palette swatches
+        swatch_height = 60
+        original_swatch = PaletteVisualizer.create_palette_image(
+            original_palette, original.width, swatch_height
+        )
+        target_swatch = PaletteVisualizer.create_palette_image(
+            target_palette, modified.width, swatch_height
+        )
+
+        # Calculate total dimensions
+        total_width = original.width + modified.width
+        total_height = max(original.height, modified.height) + swatch_height
+
+        # Create combined image
+        result = Image.new('RGB', (total_width, total_height), (255, 255, 255))
+
+        # Paste images
+        result.paste(original, (0, 0))
+        result.paste(modified, (original.width, 0))
+        result.paste(original_swatch, (0, original.height))
+        result.paste(target_swatch, (original.width, modified.height))
+
+        return result

image_generator.py

@@ -0,0 +1,100 @@
+"""
+AI Image Generation Module
+Uses Stable Diffusion to generate images from text prompts
+"""
+
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import os
+
+
+class ImageGenerator:
+    def __init__(self, model_id="stabilityai/stable-diffusion-2-1-base"):
+        """Initialize the Stable Diffusion pipeline"""
+        self.model_id = model_id
+        self.pipe = None
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    def load_model(self):
+        """Load the Stable Diffusion model"""
+        if self.pipe is None:
+            print(f"Loading model on {self.device}...")
+            self.pipe = StableDiffusionPipeline.from_pretrained(
+                self.model_id,
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
+                safety_checker=None
+            )
+            self.pipe = self.pipe.to(self.device)
+
+            # Enable memory optimizations
+            if self.device == "cuda":
+                self.pipe.enable_attention_slicing()
+
+            print("Model loaded successfully!")
+
+    def generate(
+        self,
+        prompt: str,
+        negative_prompt: str = "",
+        num_inference_steps: int = 30,
+        guidance_scale: float = 7.5,
+        width: int = 512,
+        height: int = 512,
+        seed: int = None
+    ) -> Image.Image:
+        """
+        Generate an image from a text prompt
+
+        Args:
+            prompt: Text description of desired image
+            negative_prompt: What to avoid in the image
+            num_inference_steps: Number of denoising steps (higher = better quality, slower)
+            guidance_scale: How closely to follow the prompt (7-10 recommended)
+            width: Image width (must be multiple of 8)
+            height: Image height (must be multiple of 8)
+            seed: Random seed for reproducibility
+
+        Returns:
+            PIL Image
+        """
+        self.load_model()
+
+        # Set seed for reproducibility
+        generator = None
+        if seed is not None:
+            generator = torch.Generator(device=self.device).manual_seed(seed)
+
+        # Generate image
+        with torch.inference_mode():
+            result = self.pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                width=width,
+                height=height,
+                generator=generator
+            )
+
+        return result.images[0]
+
+    def unload_model(self):
+        """Free up memory by unloading the model"""
+        if self.pipe is not None:
+            del self.pipe
+            self.pipe = None
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            print("Model unloaded")
+
+
+# Example usage
+if __name__ == "__main__":
+    generator = ImageGenerator()
+    image = generator.generate(
+        prompt="A fantasy landscape with mountains and a castle at sunset",
+        seed=42
+    )
+    image.save("test_generated.png")
+    print("Image saved as test_generated.png")

requirements.txt

@@ -0,0 +1,9 @@
+gradio==4.44.0
+diffusers==0.30.0
+transformers==4.44.0
+accelerate==0.33.0
+torch==2.4.0
+pillow==10.4.0
+numpy==1.26.4
+scikit-learn==1.5.1
+scipy==1.14.0