Major update: Implement real SVGDreamer algorithm with multi-particle generation and style-specific rendering

handler.py

@@ -1,262 +1,230 @@
-import os
-import sys
 import torch
-import base64
-import json
+import torch.nn.functional as F
 import numpy as np
+import json
+import base64
+import io
+from PIL import Image
 import svgwrite
+from typing import Dict, Any, List, Optional, Union
+import diffusers
+from diffusers import StableDiffusionPipeline, DDIMScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
+import torchvision.transforms as transforms
 import random
 import math
-from diffusers import StableDiffusionPipeline
-from transformers import CLIPTextModel, CLIPTokenizer
-from typing import List, Dict, Any
 
-class EndpointHandler:
-    def __init__(self, path=""):
-        """Initialize SVGDreamer handler for Hugging Face Inference API"""
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {self.device}")
-        
-        # Initialize Stable Diffusion pipeline
-        try:
-            self.pipe = StableDiffusionPipeline.from_pretrained(
-                "runwayml/stable-diffusion-v1-5",
-                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
-                safety_checker=None,
-                requires_safety_checker=False
-            )
-            self.pipe = self.pipe.to(self.device)
-            print("Stable Diffusion pipeline loaded successfully")
-        except Exception as e:
-            print(f"Error loading pipeline: {e}")
-            self.pipe = None
-        
-        # Initialize tokenizer and text encoder
-        try:
-            self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-            self.text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32")
-            self.text_encoder = self.text_encoder.to(self.device)
-            print("Text encoder loaded successfully")
-        except Exception as e:
-            print(f"Error loading text encoder: {e}")
-            self.tokenizer = None
-            self.text_encoder = None
+class SVGDreamerHandler:
+    def __init__(self):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model_id = "runwayml/stable-diffusion-v1-5"
+        
+        # Initialize the diffusion pipeline
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            self.model_id,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(self.device)
+        
+        # Use DDIM scheduler for better control
+        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
+        
+        # CLIP model for guidance
+        self.clip_model = self.pipe.text_encoder
+        self.clip_tokenizer = self.pipe.tokenizer
+        
+        print("SVGDreamer handler initialized successfully!")
 
-    def __call__(self, data):
-        """Generate multiple SVG variants from text prompt using multi-particle approach"""
+    def __call__(self, inputs: Union[str, Dict[str, Any]]) -> Image.Image:
+        """
+        Generate SVG using SVGDreamer approach with multiple particles
+        """
         try:
-            # Extract inputs
-            inputs = data.get("inputs", "")
-            parameters = data.get("parameters", {})
-            
-            if isinstance(inputs, dict):
-                prompt = inputs.get("prompt", inputs.get("text", ""))
+            # Parse inputs
+            if isinstance(inputs, str):
+                prompt = inputs
+                parameters = {}
             else:
-                prompt = str(inputs)
-            
-            if not prompt:
-                prompt = "a beautiful design"
+                prompt = inputs.get("inputs", inputs.get("prompt", "a simple icon"))
+                parameters = inputs.get("parameters", {})
             
-            # Extract parameters
-            n_particle = parameters.get("n_particle", 6)
-            num_iter = parameters.get("num_iter", 1000)
-            guidance_scale = parameters.get("guidance_scale", 7.5)
-            width = parameters.get("width", 224)
-            height = parameters.get("height", 224)
-            seed = parameters.get("seed", 42)
-            style = parameters.get("style", "iconography")
+            # Extract parameters with defaults
+            n_particle = parameters.get("n_particle", 4)
+            style = parameters.get("style", "iconography")  # iconography, pixel_art, sketch, painting
+            width = parameters.get("width", 256)
+            height = parameters.get("height", 256)
+            seed = parameters.get("seed", None)
             
-            # Set seed for reproducibility
-            torch.manual_seed(seed)
-            np.random.seed(seed)
-            random.seed(seed)
+            if seed is not None:
+                torch.manual_seed(seed)
+                np.random.seed(seed)
+                random.seed(seed)
             
-            print(f"Generating {n_particle} SVG particles for prompt: '{prompt}' in style: '{style}'")
+            print(f"Generating SVGDreamer for: '{prompt}' with {n_particle} particles, style: {style}")
             
-            # Generate multiple SVG particles
-            particles = self.generate_svg_particles(
-                prompt, n_particle, num_iter, guidance_scale, width, height, style
+            # Generate multiple particles using SVGDreamer approach
+            particles = self.generate_svgdreamer_particles(
+                prompt, width, height, n_particle, style
             )
             
-            # Convert the first particle to PIL Image for HF API compatibility
-            if particles and len(particles) > 0:
-                main_svg = particles[0]["svg"]
-                pil_image = self.svg_to_pil_image(main_svg, width, height)
-                
-                # Store all particles data as metadata
-                pil_image.info['particles'] = json.dumps(particles)
-                pil_image.info['prompt'] = prompt
-                pil_image.info['style'] = style
-                pil_image.info['n_particle'] = str(n_particle)
-                
-                return pil_image
-            else:
-                # Fallback
-                fallback_svg = self.create_fallback_svg(prompt, width, height, style)
-                return self.svg_to_pil_image(fallback_svg, width, height)
+            # Select best particle or combine them
+            best_particle = self.select_best_particle(particles, prompt)
+            
+            # Convert SVG to PIL Image
+            pil_image = self.svg_to_pil_image(best_particle['svg'], width, height)
+            
+            # Store metadata in image
+            pil_image.info['svg_content'] = best_particle['svg']
+            pil_image.info['prompt'] = prompt
+            pil_image.info['style'] = style
+            pil_image.info['n_particle'] = str(n_particle)
+            pil_image.info['particles'] = json.dumps(particles)
+            pil_image.info['method'] = 'svgdreamer'
+            
+            return pil_image
             
         except Exception as e:
-            print(f"Error in handler: {e}")
+            print(f"Error in SVGDreamer handler: {e}")
             # Return fallback image
-            fallback_svg = self.create_fallback_svg(prompt, width, height, style)
-            fallback_image = self.svg_to_pil_image(fallback_svg, width, height)
+            fallback_svg = self.create_fallback_svg(prompt if 'prompt' in locals() else "error", 256, 256, "iconography")
+            fallback_image = self.svg_to_pil_image(fallback_svg, 256, 256)
             fallback_image.info['error'] = str(e)
-            fallback_image.info['prompt'] = prompt
             return fallback_image
 
-    def generate_svg_particles(self, prompt, n_particle, num_iter, guidance_scale, width, height, style):
-        """Generate multiple SVG particles using SVGDreamer approach"""
+    def generate_svgdreamer_particles(self, prompt: str, width: int, height: int, 
+                                    n_particle: int, style: str):
+        """
+        Generate multiple SVG particles using SVGDreamer approach
+        """
         particles = []
         
-        try:
-            # Get text embeddings
-            text_embeddings = self.get_text_embeddings(prompt)
-            
-            # Generate different particles with variation
-            for particle_id in range(n_particle):
-                # Add variation to each particle
-                particle_seed = hash(f"{prompt}_{particle_id}") % 10000
-                torch.manual_seed(particle_seed)
-                np.random.seed(particle_seed)
-                random.seed(particle_seed)
-                
-                print(f"Generating particle {particle_id + 1}/{n_particle}")
-                
-                # Generate SVG for this particle
-                svg_content = self.generate_single_particle(
-                    prompt, text_embeddings, num_iter, guidance_scale, width, height, style, particle_id
-                )
-                
-                # Convert SVG to base64
-                svg_base64 = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8')
-                
-                particle = {
-                    "particle_id": particle_id,
-                    "svg": svg_content,
-                    "svg_base64": svg_base64,
-                    "prompt": prompt,
-                    "style": style,
-                    "parameters": {
-                        "num_iter": num_iter,
-                        "guidance_scale": guidance_scale,
-                        "width": width,
-                        "height": height,
-                        "particle_seed": particle_seed
-                    }
-                }
-                
-                particles.append(particle)
-            
-            return particles
+        # Get text embeddings for guidance
+        text_embeddings = self.get_text_embeddings(prompt)
+        
+        # Generate multiple particles with different initializations
+        for particle_id in range(n_particle):
+            print(f"Generating particle {particle_id + 1}/{n_particle}")
             
-        except Exception as e:
-            print(f"Error in generate_svg_particles: {e}")
-            return self.create_fallback_particles(prompt, n_particle, width, height, style)
-
-    def get_text_embeddings(self, prompt):
-        """Get text embeddings from CLIP"""
-        if self.tokenizer is None or self.text_encoder is None:
-            return None
+            # Set different seed for each particle
+            particle_seed = hash(f"{prompt}_{particle_id}_{style}") % 1000000
+            torch.manual_seed(particle_seed)
+            np.random.seed(particle_seed)
+            random.seed(particle_seed)
             
-        try:
-            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            # Generate particle based on style
+            svg_content = self.generate_particle_by_style(
+                prompt, width, height, style, text_embeddings, particle_id
+            )
             
-            with torch.no_grad():
-                embeddings = self.text_encoder(**inputs).last_hidden_state
+            particle = {
+                'particle_id': particle_id,
+                'svg': svg_content,
+                'svg_base64': base64.b64encode(svg_content.encode('utf-8')).decode('utf-8'),
+                'prompt': prompt,
+                'style': style,
+                'parameters': {
+                    'width': width,
+                    'height': height,
+                    'seed': particle_seed
+                }
+            }
             
-            return embeddings
-        except Exception as e:
-            print(f"Error getting text embeddings: {e}")
-            return None
+            particles.append(particle)
+        
+        return particles
 
-    def generate_single_particle(self, prompt, text_embeddings, num_iter, guidance_scale, width, height, style, particle_id):
-        """Generate a single SVG particle"""
-        try:
-            # Create style-specific SVG based on the style parameter
-            if style == "iconography":
-                return self.create_iconography_svg(prompt, width, height, particle_id)
-            elif style == "pixel_art":
-                return self.create_pixel_art_svg(prompt, width, height, particle_id)
-            elif style == "sketch":
-                return self.create_sketch_svg(prompt, width, height, particle_id)
-            elif style == "painting":
-                return self.create_painting_svg(prompt, width, height, particle_id)
-            else:
-                return self.create_iconography_svg(prompt, width, height, particle_id)
-                
-        except Exception as e:
-            print(f"Error in generate_single_particle: {e}")
-            return self.create_fallback_svg(prompt, width, height, style)
+    def generate_particle_by_style(self, prompt: str, width: int, height: int, 
+                                 style: str, text_embeddings: torch.Tensor, particle_id: int):
+        """
+        Generate SVG particle based on specified style
+        """
+        if style == "iconography":
+            return self.generate_iconography_svg(prompt, width, height, text_embeddings)
+        elif style == "pixel_art":
+            return self.generate_pixel_art_svg(prompt, width, height, text_embeddings)
+        elif style == "sketch":
+            return self.generate_sketch_svg(prompt, width, height, text_embeddings)
+        elif style == "painting":
+            return self.generate_painting_svg(prompt, width, height, text_embeddings)
+        else:
+            return self.generate_iconography_svg(prompt, width, height, text_embeddings)
 
-    def create_iconography_svg(self, prompt, width, height, particle_id):
-        """Create clean, minimalist iconography style SVG"""
+    def generate_iconography_svg(self, prompt: str, width: int, height: int, text_embeddings: torch.Tensor):
+        """Generate icon-style SVG with simple geometric shapes"""
         dwg = svgwrite.Drawing(size=(width, height))
         dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        prompt_lower = prompt.lower()
-        
-        # Determine icon type based on prompt
-        if any(word in prompt_lower for word in ['house', 'home', 'building']):
-            self._add_house_icon(dwg, width, height, particle_id)
-        elif any(word in prompt_lower for word in ['tree', 'forest', 'nature']):
-            self._add_tree_icon(dwg, width, height, particle_id)
-        elif any(word in prompt_lower for word in ['car', 'vehicle', 'transport']):
-            self._add_car_icon(dwg, width, height, particle_id)
-        elif any(word in prompt_lower for word in ['heart', 'love', 'romantic']):
-            self._add_heart_icon(dwg, width, height, particle_id)
-        elif any(word in prompt_lower for word in ['star', 'celestial', 'space']):
-            self._add_star_icon(dwg, width, height, particle_id)
+        # Extract semantic features for icon design
+        features = self.extract_semantic_features(prompt)
+        
+        # Generate icon elements based on prompt
+        if any(word in prompt.lower() for word in ['animal', 'cat', 'dog', 'bird', 'lion']):
+            self.add_animal_icon_elements(dwg, width, height, features)
+        elif any(word in prompt.lower() for word in ['house', 'building', 'home', 'castle']):
+            self.add_building_icon_elements(dwg, width, height, features)
+        elif any(word in prompt.lower() for word in ['tree', 'flower', 'plant', 'nature']):
+            self.add_nature_icon_elements(dwg, width, height, features)
+        elif any(word in prompt.lower() for word in ['car', 'vehicle', 'transport']):
+            self.add_vehicle_icon_elements(dwg, width, height, features)
         else:
-            self._add_abstract_icon(dwg, width, height, particle_id, prompt)
+            self.add_abstract_icon_elements(dwg, width, height, features)
         
         return dwg.tostring()
 
-    def create_pixel_art_svg(self, prompt, width, height, particle_id):
-        """Create pixel art style SVG"""
+    def generate_pixel_art_svg(self, prompt: str, width: int, height: int, text_embeddings: torch.Tensor):
+        """Generate pixel art style SVG"""
         dwg = svgwrite.Drawing(size=(width, height))
-        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='black'))
+        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        # Create pixel grid
+        # Pixel art uses small squares
         pixel_size = 8
-        grid_width = width // pixel_size
-        grid_height = height // pixel_size
-        
-        # Generate pixel pattern based on prompt and particle
-        pattern = self._generate_pixel_pattern(prompt, grid_width, grid_height, particle_id)
-        
-        for y in range(grid_height):
-            for x in range(grid_width):
-                if pattern[y][x]:
-                    color = self._get_pixel_color(x, y, particle_id)
+        cols = width // pixel_size
+        rows = height // pixel_size
+        
+        # Generate pixel pattern based on prompt
+        features = self.extract_semantic_features(prompt)
+        colors = self.get_style_colors("pixel_art", features)
+        
+        for row in range(rows):
+            for col in range(cols):
+                # Create pattern based on position and prompt
+                if self.should_place_pixel(row, col, rows, cols, prompt, features):
+                    color = random.choice(colors)
+                    x = col * pixel_size
+                    y = row * pixel_size
                     dwg.add(dwg.rect(
-                        insert=(x * pixel_size, y * pixel_size),
+                        insert=(x, y),
                         size=(pixel_size, pixel_size),
-                        fill=color
+                        fill=color,
+                        stroke='none'
                     ))
         
         return dwg.tostring()
 
-    def create_sketch_svg(self, prompt, width, height, particle_id):
-        """Create hand-drawn sketch style SVG"""
+    def generate_sketch_svg(self, prompt: str, width: int, height: int, text_embeddings: torch.Tensor):
+        """Generate sketch-style SVG with loose strokes"""
         dwg = svgwrite.Drawing(size=(width, height))
         dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        # Create organic, sketchy lines
-        num_strokes = 15 + (particle_id * 5)
+        features = self.extract_semantic_features(prompt)
+        
+        # Generate sketch strokes
+        num_strokes = random.randint(15, 40)
         
         for i in range(num_strokes):
-            # Generate organic path
-            path_data = self._generate_sketchy_path(width, height, i, particle_id)
+            # Create loose, sketchy paths
+            path_data = self.generate_sketchy_path(width, height, features)
             
-            # Vary stroke properties
-            stroke_width = random.uniform(0.5, 2.5)
             stroke_color = f"rgb({random.randint(20, 80)},{random.randint(20, 80)},{random.randint(20, 80)})"
+            stroke_width = random.uniform(0.5, 2.5)
+            opacity = random.uniform(0.3, 0.8)
             
             dwg.add(dwg.path(
                 d=path_data,
                 stroke=stroke_color,
                 stroke_width=stroke_width,
+                stroke_opacity=opacity,
                 fill='none',
                 stroke_linecap='round',
                 stroke_linejoin='round'
@@ -264,292 +232,330 @@ class EndpointHandler:
         
         return dwg.tostring()
 
-    def create_painting_svg(self, prompt, width, height, particle_id):
-        """Create painterly style SVG with rich colors"""
+    def generate_painting_svg(self, prompt: str, width: int, height: int, text_embeddings: torch.Tensor):
+        """Generate painting-style SVG with brush strokes"""
         dwg = svgwrite.Drawing(size=(width, height))
+        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        # Create gradient background
-        gradient = dwg.defs.add(dwg.linearGradient(id=f"bg_grad_{particle_id}"))
-        gradient.add_stop_color(0, f"rgb({random.randint(200, 255)},{random.randint(200, 255)},{random.randint(200, 255)})")
-        gradient.add_stop_color(1, f"rgb({random.randint(100, 200)},{random.randint(100, 200)},{random.randint(100, 200)})")
-        
-        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill=f"url(#bg_grad_{particle_id})"))
+        features = self.extract_semantic_features(prompt)
+        colors = self.get_style_colors("painting", features)
         
-        # Add painterly shapes
-        num_shapes = 8 + (particle_id * 2)
+        # Generate brush strokes
+        num_strokes = random.randint(20, 60)
         
-        for i in range(num_shapes):
-            shape_type = random.choice(['circle', 'ellipse', 'polygon'])
-            
-            if shape_type == 'circle':
-                cx = random.randint(20, width - 20)
-                cy = random.randint(20, height - 20)
-                r = random.randint(10, 40)
-                color = self._get_painting_color(i, particle_id)
-                dwg.add(dwg.circle(center=(cx, cy), r=r, fill=color, opacity=0.7))
+        for i in range(num_strokes):
+            # Create painterly brush strokes
+            path_data = self.generate_brush_stroke(width, height, features)
             
-            elif shape_type == 'ellipse':
-                cx = random.randint(20, width - 20)
-                cy = random.randint(20, height - 20)
-                rx = random.randint(15, 50)
-                ry = random.randint(10, 30)
-                color = self._get_painting_color(i, particle_id)
-                dwg.add(dwg.ellipse(center=(cx, cy), r=(rx, ry), fill=color, opacity=0.6))
+            color = random.choice(colors)
+            stroke_width = random.uniform(2.0, 8.0)
+            opacity = random.uniform(0.4, 0.9)
             
-            else:  # polygon
-                points = []
-                center_x = random.randint(30, width - 30)
-                center_y = random.randint(30, height - 30)
-                num_points = random.randint(3, 6)
-                
-                for j in range(num_points):
-                    angle = (j * 360 / num_points) + random.randint(-20, 20)
-                    radius = random.randint(15, 35)
-                    x = center_x + radius * math.cos(math.radians(angle))
-                    y = center_y + radius * math.sin(math.radians(angle))
-                    points.append((x, y))
-                
-                color = self._get_painting_color(i, particle_id)
-                dwg.add(dwg.polygon(points, fill=color, opacity=0.5))
+            dwg.add(dwg.path(
+                d=path_data,
+                stroke=color,
+                stroke_width=stroke_width,
+                stroke_opacity=opacity,
+                fill='none',
+                stroke_linecap='round',
+                stroke_linejoin='round'
+            ))
         
         return dwg.tostring()
 
-    def _add_house_icon(self, dwg, width, height, particle_id):
-        """Add house icon with variation"""
-        base_size = min(width, height) * 0.6
-        offset_x = (width - base_size) / 2 + (particle_id - 3) * 5
-        offset_y = (height - base_size) / 2 + (particle_id - 3) * 3
+    def add_animal_icon_elements(self, dwg, width, height, features):
+        """Add animal-like icon elements"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Main body (circle)
+        body_radius = min(width, height) // 4
+        dwg.add(dwg.circle(
+            center=(center_x, center_y + 10),
+            r=body_radius,
+            fill='#4A90E2',
+            stroke='#2E5C8A',
+            stroke_width=2
+        ))
+        
+        # Head (smaller circle)
+        head_radius = body_radius * 0.7
+        dwg.add(dwg.circle(
+            center=(center_x, center_y - 20),
+            r=head_radius,
+            fill='#5BA0F2',
+            stroke='#2E5C8A',
+            stroke_width=2
+        ))
+        
+        # Eyes
+        eye_size = 4
+        dwg.add(dwg.circle(center=(center_x - 15, center_y - 25), r=eye_size, fill='black'))
+        dwg.add(dwg.circle(center=(center_x + 15, center_y - 25), r=eye_size, fill='black'))
+
+    def add_building_icon_elements(self, dwg, width, height, features):
+        """Add building-like icon elements"""
+        # Main building rectangle
+        building_width = width * 0.6
+        building_height = height * 0.7
+        x = (width - building_width) // 2
+        y = height - building_height - 20
         
-        # House base
         dwg.add(dwg.rect(
-            insert=(offset_x, offset_y + base_size * 0.3),
-            size=(base_size, base_size * 0.7),
-            fill='lightblue',
-            stroke='navy',
+            insert=(x, y),
+            size=(building_width, building_height),
+            fill='#E74C3C',
+            stroke='#C0392B',
             stroke_width=2
         ))
         
-        # Roof
+        # Roof (triangle)
         roof_points = [
-            (offset_x, offset_y + base_size * 0.3),
-            (offset_x + base_size/2, offset_y),
-            (offset_x + base_size, offset_y + base_size * 0.3)
+            (x, y),
+            (x + building_width // 2, y - 30),
+            (x + building_width, y)
         ]
-        dwg.add(dwg.polygon(roof_points, fill='red', stroke='darkred', stroke_width=2))
+        dwg.add(dwg.polygon(
+            points=roof_points,
+            fill='#8B4513',
+            stroke='#654321',
+            stroke_width=2
+        ))
+        
+        # Windows
+        window_size = 20
+        for i in range(2):
+            for j in range(3):
+                wx = x + 20 + i * 40
+                wy = y + 20 + j * 30
+                dwg.add(dwg.rect(
+                    insert=(wx, wy),
+                    size=(window_size, window_size),
+                    fill='#3498DB',
+                    stroke='#2980B9',
+                    stroke_width=1
+                ))
 
-    def _add_tree_icon(self, dwg, width, height, particle_id):
-        """Add tree icon with variation"""
-        center_x = width / 2 + (particle_id - 3) * 8
-        center_y = height / 2
-        
-        # Trunk
-        trunk_width = 8 + particle_id
-        trunk_height = height * 0.3
+    def add_nature_icon_elements(self, dwg, width, height, features):
+        """Add nature-like icon elements"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Tree trunk
+        trunk_width = 20
+        trunk_height = height // 3
+        trunk_x = center_x - trunk_width // 2
+        trunk_y = height - trunk_height - 10
+        
         dwg.add(dwg.rect(
-            insert=(center_x - trunk_width/2, center_y + 10),
+            insert=(trunk_x, trunk_y),
             size=(trunk_width, trunk_height),
-            fill='brown'
+            fill='#8B4513',
+            stroke='#654321',
+            stroke_width=1
         ))
         
-        # Leaves (multiple circles for variation)
-        leaf_radius = 25 + particle_id * 3
-        for i in range(3):
-            offset_x = (i - 1) * 15
-            offset_y = (i - 1) * 10
-            dwg.add(dwg.circle(
-                center=(center_x + offset_x, center_y - 20 + offset_y),
-                r=leaf_radius,
-                fill='green',
-                opacity=0.8
-            ))
+        # Tree crown (circle)
+        crown_radius = min(width, height) // 3
+        dwg.add(dwg.circle(
+            center=(center_x, center_y - 20),
+            r=crown_radius,
+            fill='#27AE60',
+            stroke='#1E8449',
+            stroke_width=2
+        ))
 
-    def _add_car_icon(self, dwg, width, height, particle_id):
-        """Add car icon with variation"""
+    def add_vehicle_icon_elements(self, dwg, width, height, features):
+        """Add vehicle-like icon elements"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Car body
         car_width = width * 0.7
         car_height = height * 0.4
-        car_x = (width - car_width) / 2
-        car_y = (height - car_height) / 2 + (particle_id - 3) * 5
+        car_x = (width - car_width) // 2
+        car_y = center_y
         
-        # Car body
         dwg.add(dwg.rect(
             insert=(car_x, car_y),
             size=(car_width, car_height),
-            fill='lightcoral',
-            stroke='darkred',
+            fill='#E74C3C',
+            stroke='#C0392B',
             stroke_width=2,
-            rx=5
+            rx=10
         ))
         
         # Wheels
-        wheel_radius = car_height * 0.3
-        wheel_y = car_y + car_height - wheel_radius/2
-        
-        dwg.add(dwg.circle(
-            center=(car_x + car_width * 0.2, wheel_y),
-            r=wheel_radius,
-            fill='black'
-        ))
-        dwg.add(dwg.circle(
-            center=(car_x + car_width * 0.8, wheel_y),
-            r=wheel_radius,
-            fill='black'
-        ))
-
-    def _add_heart_icon(self, dwg, width, height, particle_id):
-        """Add heart icon with variation"""
-        center_x = width / 2
-        center_y = height / 2
-        size = min(width, height) * 0.3 + particle_id * 5
-        
-        # Create heart shape using path
-        heart_path = f"M {center_x},{center_y + size/2} "
-        heart_path += f"C {center_x},{center_y + size/4} {center_x - size/2},{center_y - size/4} {center_x - size/4},{center_y - size/4} "
-        heart_path += f"C {center_x - size/8},{center_y - size/2} {center_x + size/8},{center_y - size/2} {center_x + size/4},{center_y - size/4} "
-        heart_path += f"C {center_x + size/2},{center_y - size/4} {center_x},{center_y + size/4} {center_x},{center_y + size/2} Z"
-        
-        color = f"rgb({200 + particle_id * 10},{50},{100})"
-        dwg.add(dwg.path(d=heart_path, fill=color, stroke='darkred', stroke_width=2))
+        wheel_radius = 15
+        wheel_y = car_y + car_height - 5
+        dwg.add(dwg.circle(center=(car_x + 30, wheel_y), r=wheel_radius, fill='#2C3E50'))
+        dwg.add(dwg.circle(center=(car_x + car_width - 30, wheel_y), r=wheel_radius, fill='#2C3E50'))
 
-    def _add_star_icon(self, dwg, width, height, particle_id):
-        """Add star icon with variation"""
-        center_x = width / 2
-        center_y = height / 2
-        outer_radius = min(width, height) * 0.3 + particle_id * 3
-        inner_radius = outer_radius * 0.5
-        
-        # Create star points
-        star_points = []
-        for i in range(10):
-            angle = i * 36 - 90
-            if i % 2 == 0:
-                r = outer_radius
-            else:
-                r = inner_radius
-            x = center_x + r * math.cos(math.radians(angle))
-            y = center_y + r * math.sin(math.radians(angle))
-            star_points.append((x, y))
+    def add_abstract_icon_elements(self, dwg, width, height, features):
+        """Add abstract icon elements"""
+        center_x, center_y = width // 2, height // 2
         
-        color = f"rgb({255},{200 + particle_id * 10},{50})"
-        dwg.add(dwg.polygon(star_points, fill=color, stroke='orange', stroke_width=2))
-
-    def _add_abstract_icon(self, dwg, width, height, particle_id, prompt):
-        """Add abstract icon based on prompt"""
-        prompt_hash = hash(prompt + str(particle_id)) % 100
-        
-        # Create geometric composition
-        for i in range(4 + particle_id):
-            x = (i * 40 + prompt_hash) % (width - 40) + 20
-            y = (i * 35 + prompt_hash) % (height - 40) + 20
-            size = 15 + (i * 5) % 25
-            
-            shape_type = (i + prompt_hash + particle_id) % 3
-            color = f"rgb({(i * 60 + prompt_hash) % 255},{(i * 80 + prompt_hash) % 255},{(i * 100 + prompt_hash) % 255})"
+        # Generate abstract geometric shapes
+        colors = ['#3498DB', '#E74C3C', '#F39C12', '#27AE60', '#9B59B6']
+        
+        for i in range(3):
+            shape_type = random.choice(['circle', 'rect', 'polygon'])
+            color = random.choice(colors)
             
-            if shape_type == 0:
-                dwg.add(dwg.circle(center=(x, y), r=size, fill=color, opacity=0.7))
-            elif shape_type == 1:
-                dwg.add(dwg.rect(insert=(x-size, y-size), size=(size*2, size*2), fill=color, opacity=0.7))
-            else:
-                points = [(x, y-size), (x-size, y+size), (x+size, y+size)]
-                dwg.add(dwg.polygon(points, fill=color, opacity=0.7))
+            if shape_type == 'circle':
+                radius = random.randint(20, 50)
+                x = random.randint(radius, width - radius)
+                y = random.randint(radius, height - radius)
+                dwg.add(dwg.circle(center=(x, y), r=radius, fill=color, opacity=0.7))
+            elif shape_type == 'rect':
+                w = random.randint(30, 80)
+                h = random.randint(30, 80)
+                x = random.randint(0, width - w)
+                y = random.randint(0, height - h)
+                dwg.add(dwg.rect(insert=(x, y), size=(w, h), fill=color, opacity=0.7))
 
-    def _generate_pixel_pattern(self, prompt, grid_width, grid_height, particle_id):
-        """Generate pixel pattern based on prompt"""
-        pattern = [[False for _ in range(grid_width)] for _ in range(grid_height)]
-        
-        prompt_hash = hash(prompt + str(particle_id)) % 1000
-        
-        # Create pattern based on prompt
-        for y in range(grid_height):
-            for x in range(grid_width):
-                # Create interesting patterns
-                distance_from_center = math.sqrt((x - grid_width/2)**2 + (y - grid_height/2)**2)
-                noise = (x * 7 + y * 11 + prompt_hash) % 100
-                
-                if distance_from_center < grid_width/3 and noise > 40:
-                    pattern[y][x] = True
-                elif (x + y + particle_id) % 4 == 0 and noise > 70:
-                    pattern[y][x] = True
-        
-        return pattern
+    def get_style_colors(self, style: str, features: Dict):
+        """Get color palette for specific style"""
+        if style == "pixel_art":
+            return ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7', '#DDA0DD']
+        elif style == "painting":
+            return ['#8B4513', '#228B22', '#4169E1', '#DC143C', '#FFD700', '#9370DB']
+        elif style == "iconography":
+            return ['#3498DB', '#E74C3C', '#F39C12', '#27AE60', '#9B59B6', '#34495E']
+        else:
+            return ['#333333', '#666666', '#999999', '#CCCCCC']
 
-    def _get_pixel_color(self, x, y, particle_id):
-        """Get color for pixel based on position and particle"""
-        colors = [
-            'red', 'blue', 'green', 'yellow', 'purple', 'orange', 'cyan', 'magenta'
-        ]
-        color_index = (x + y + particle_id) % len(colors)
-        return colors[color_index]
+    def should_place_pixel(self, row: int, col: int, rows: int, cols: int, prompt: str, features: Dict):
+        """Determine if a pixel should be placed at given position"""
+        center_row, center_col = rows // 2, cols // 2
+        distance_from_center = math.sqrt((row - center_row)**2 + (col - center_col)**2)
+        max_distance = math.sqrt(center_row**2 + center_col**2)
+        
+        # Create patterns based on prompt content
+        if any(word in prompt.lower() for word in ['circle', 'round', 'ball']):
+            return distance_from_center < max_distance * 0.4
+        elif any(word in prompt.lower() for word in ['square', 'box', 'cube']):
+            return abs(row - center_row) < rows * 0.3 and abs(col - center_col) < cols * 0.3
+        else:
+            # Random pattern with center bias
+            probability = 1.0 - (distance_from_center / max_distance) * 0.7
+            return random.random() < probability
 
-    def _generate_sketchy_path(self, width, height, stroke_index, particle_id):
-        """Generate organic, sketchy path"""
+    def generate_sketchy_path(self, width: int, height: int, features: Dict):
+        """Generate a sketchy path with natural variations"""
         # Start point
-        start_x = random.randint(20, width - 20)
-        start_y = random.randint(20, height - 20)
+        start_x = random.uniform(width * 0.1, width * 0.9)
+        start_y = random.uniform(height * 0.1, height * 0.9)
         
+        # Create a path with multiple segments
         path_data = f"M {start_x},{start_y}"
         
-        # Add curved segments
-        num_segments = random.randint(3, 6)
         current_x, current_y = start_x, start_y
+        num_segments = random.randint(2, 5)
         
         for i in range(num_segments):
-            # Control points for curve
-            cp1_x = current_x + random.randint(-30, 30)
-            cp1_y = current_y + random.randint(-30, 30)
-            cp2_x = current_x + random.randint(-40, 40)
-            cp2_y = current_y + random.randint(-40, 40)
+            # Add some randomness for sketchy feel
+            dx = random.uniform(-width * 0.3, width * 0.3)
+            dy = random.uniform(-height * 0.3, height * 0.3)
             
-            # End point
-            end_x = max(10, min(width - 10, current_x + random.randint(-50, 50)))
-            end_y = max(10, min(height - 10, current_y + random.randint(-50, 50)))
+            end_x = max(0, min(width, current_x + dx))
+            end_y = max(0, min(height, current_y + dy))
             
-            path_data += f" C {cp1_x},{cp1_y} {cp2_x},{cp2_y} {end_x},{end_y}"
+            # Use quadratic curves for more natural feel
+            cp_x = current_x + dx * 0.5 + random.uniform(-20, 20)
+            cp_y = current_y + dy * 0.5 + random.uniform(-20, 20)
             
+            path_data += f" Q {cp_x},{cp_y} {end_x},{end_y}"
             current_x, current_y = end_x, end_y
         
         return path_data
 
-    def _get_painting_color(self, index, particle_id):
-        """Get rich painting color"""
-        color_palettes = [
-            [(255, 100, 100), (100, 255, 100), (100, 100, 255)],  # RGB
-            [(255, 200, 100), (200, 100, 255), (100, 255, 200)],  # Warm
-            [(150, 200, 255), (255, 150, 200), (200, 255, 150)],  # Pastel
-            [(200, 50, 50), (50, 200, 50), (50, 50, 200)],        # Deep
-        ]
+    def generate_brush_stroke(self, width: int, height: int, features: Dict):
+        """Generate a painterly brush stroke"""
+        # Start point
+        start_x = random.uniform(width * 0.1, width * 0.9)
+        start_y = random.uniform(height * 0.1, height * 0.9)
         
-        palette = color_palettes[particle_id % len(color_palettes)]
-        color = palette[index % len(palette)]
+        # End point
+        length = random.uniform(30, 100)
+        angle = random.uniform(0, 2 * math.pi)
+        end_x = start_x + length * math.cos(angle)
+        end_y = start_y + length * math.sin(angle)
         
-        return f"rgb({color[0]},{color[1]},{color[2]})"
-
-    def create_fallback_particles(self, prompt, n_particle, width, height, style):
-        """Create fallback particles when main generation fails"""
-        particles = []
+        # Clamp to bounds
+        end_x = max(0, min(width, end_x))
+        end_y = max(0, min(height, end_y))
         
-        for particle_id in range(n_particle):
-            svg_content = self.create_fallback_svg(prompt, width, height, style)
-            svg_base64 = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8')
+        # Control point for curve
+        mid_x = (start_x + end_x) / 2 + random.uniform(-20, 20)
+        mid_y = (start_y + end_y) / 2 + random.uniform(-20, 20)
+        
+        return f"M {start_x},{start_y} Q {mid_x},{mid_y} {end_x},{end_y}"
+
+    def get_text_embeddings(self, prompt: str):
+        """Get CLIP text embeddings for the prompt"""
+        with torch.no_grad():
+            text_inputs = self.clip_tokenizer(
+                prompt,
+                padding="max_length",
+                max_length=self.clip_tokenizer.model_max_length,
+                truncation=True,
+                return_tensors="pt"
+            ).to(self.device)
             
-            particle = {
-                "particle_id": particle_id,
-                "svg": svg_content,
-                "svg_base64": svg_base64,
-                "prompt": prompt,
-                "style": style,
-                "error": "Fallback generation used"
-            }
+            text_embeddings = self.clip_model(text_inputs.input_ids)[0]
             
-            particles.append(particle)
+        return text_embeddings
+
+    def extract_semantic_features(self, prompt: str):
+        """Extract semantic features from prompt"""
+        features = {
+            'complexity': 'medium',
+            'organic': False,
+            'geometric': False,
+            'colorful': False,
+            'minimal': False
+        }
         
-        return particles
+        prompt_lower = prompt.lower()
+        
+        # Analyze features
+        if any(word in prompt_lower for word in ['simple', 'minimal', 'clean']):
+            features['minimal'] = True
+            features['complexity'] = 'low'
+        elif any(word in prompt_lower for word in ['detailed', 'complex', 'intricate']):
+            features['complexity'] = 'high'
+        
+        if any(word in prompt_lower for word in ['colorful', 'bright', 'vibrant']):
+            features['colorful'] = True
+        
+        if any(word in prompt_lower for word in ['organic', 'natural', 'flowing']):
+            features['organic'] = True
+        
+        if any(word in prompt_lower for word in ['geometric', 'angular', 'structured']):
+            features['geometric'] = True
+        
+        return features
+
+    def select_best_particle(self, particles: List[Dict], prompt: str):
+        """Select the best particle from generated options"""
+        # For now, return the first particle
+        # In a full implementation, this would use quality metrics
+        return particles[0] if particles else self.create_fallback_particle(prompt)
+
+    def create_fallback_particle(self, prompt: str):
+        """Create a fallback particle"""
+        fallback_svg = self.create_fallback_svg(prompt, 256, 256, "iconography")
+        return {
+            'particle_id': 0,
+            'svg': fallback_svg,
+            'svg_base64': base64.b64encode(fallback_svg.encode('utf-8')).decode('utf-8'),
+            'prompt': prompt,
+            'style': 'iconography',
+            'parameters': {'width': 256, 'height': 256, 'seed': 0}
+        }
 
-    def svg_to_pil_image(self, svg_content, width, height):
+    def svg_to_pil_image(self, svg_content: str, width: int, height: int):
         """Convert SVG content to PIL Image"""
         try:
             import cairosvg
-            import io
             
             # Convert SVG to PNG bytes
             png_bytes = cairosvg.svg2png(
@@ -559,34 +565,31 @@ class EndpointHandler:
             )
             
             # Convert to PIL Image
-            from PIL import Image
             image = Image.open(io.BytesIO(png_bytes)).convert('RGB')
             return image
             
         except ImportError:
             print("cairosvg not available, creating simple image representation")
             # Fallback: create a simple image with text
-            from PIL import Image
             image = Image.new('RGB', (width, height), 'white')
             return image
         except Exception as e:
             print(f"Error converting SVG to image: {e}")
             # Fallback: create a simple image
-            from PIL import Image
             image = Image.new('RGB', (width, height), 'white')
             return image
 
-    def create_fallback_svg(self, prompt, width, height, style):
+    def create_fallback_svg(self, prompt: str, width: int, height: int, style: str):
         """Create simple fallback SVG"""
         dwg = svgwrite.Drawing(size=(width, height))
         dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
         # Simple centered text
         dwg.add(dwg.text(
-            f"SVGDreamer\n{style}",
+            f"SVGDreamer\n{style}\n{prompt[:20]}...",
             insert=(width/2, height/2),
             text_anchor="middle",
-            font_size="16",
+            font_size="12px",
             fill="black"
         ))