Upload folder using huggingface_hub

handler.py

@@ -67,7 +67,7 @@ class SVGDreamerHandler:
         return Args()
     
     def __call__(self, data: Dict[str, Any]):
-        """Process input and return PIL Image"""
+        """Generate multi-particle SVG from text prompt"""
         try:
             # Extract inputs
             if isinstance(data, dict):
@@ -81,6 +81,7 @@ class SVGDreamerHandler:
                 prompt = "a simple drawing"
             
             # Extract parameters
+            n_particle = parameters.get("n_particle", 6)
             width = parameters.get("width", 224)
             height = parameters.get("height", 224)
             seed = parameters.get("seed", 42)
@@ -89,57 +90,195 @@ class SVGDreamerHandler:
             # Set random seed
             np.random.seed(seed)
             
-            # Create PIL Image for proper serialization
-            from PIL import Image, ImageDraw
+            # Generate the best particle SVG (simulate particle selection)
+            best_svg = self._generate_particle_svg(prompt, width, height, style, seed)
             
-            img = Image.new('RGB', (width, height), 'white')
-            draw = ImageDraw.Draw(img)
+            return best_svg
             
-            # Different color schemes based on style
-            if style == "iconography":
-                colors = [(44, 62, 80), (231, 76, 60), (52, 152, 219), (46, 204, 113)]
-            elif style == "pixel_art":
-                colors = [(255, 107, 107), (78, 205, 196), (69, 183, 209), (150, 206, 180)]
-            else:
-                colors = [(52, 73, 94), (230, 126, 34), (26, 188, 156), (142, 68, 173)]
-            
-            # Generate patterns based on style
-            if style == "iconography":
-                # Clean geometric shapes
-                for i in range(6):
-                    x = np.random.randint(20, width-40)
-                    y = np.random.randint(20, height-40)
-                    size = np.random.randint(15, 35)
-                    color = colors[i % len(colors)]
-                    draw.ellipse([x, y, x+size, y+size], fill=color)
-            
-            elif style == "pixel_art":
-                # Pixelated squares
-                for i in range(12):
-                    x = np.random.randint(0, width-20)
-                    y = np.random.randint(0, height-20)
-                    size = 15
-                    color = colors[i % len(colors)]
-                    draw.rectangle([x, y, x+size, y+size], fill=color)
+        except Exception as e:
+            # Return error SVG
+            return f'<svg width="224" height="224" xmlns="http://www.w3.org/2000/svg"><text x="10" y="20" fill="red">Error: {str(e)}</text></svg>'
+    
+    def _generate_particle_svg(self, prompt: str, width: int, height: int, style: str, seed: int) -> str:
+        """Generate SVG using particle-based optimization simulation"""
+        svg_header = f'<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}">'
+        svg_footer = '</svg>'
+        
+        paths = []
+        prompt_lower = prompt.lower()
+        
+        # Style-based color palettes
+        if style == "iconography":
+            colors = ["#2C3E50", "#E74C3C", "#3498DB", "#2ECC71", "#F39C12", "#9B59B6"]
+        elif style == "pixel_art":
+            colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD"]
+        elif style == "sketch":
+            colors = ["#34495E", "#7F8C8D", "#95A5A6", "#BDC3C7", "#ECF0F1", "#2C3E50"]
+        else:  # painting
+            colors = ["#8E44AD", "#3498DB", "#E67E22", "#E74C3C", "#F1C40F", "#27AE60"]
+        
+        # Generate content based on prompt analysis
+        if any(word in prompt_lower for word in ['icon', 'logo', 'symbol', 'simple']):
+            self._add_icon_elements(paths, width, height, colors, style)
+        elif any(word in prompt_lower for word in ['landscape', 'mountain', 'nature', 'scene']):
+            self._add_landscape_elements(paths, width, height, colors, style)
+        elif any(word in prompt_lower for word in ['character', 'person', 'face', 'figure']):
+            self._add_character_elements(paths, width, height, colors, style)
+        elif any(word in prompt_lower for word in ['abstract', 'pattern', 'geometric']):
+            self._add_abstract_elements(paths, width, height, colors, style)
+        else:
+            self._add_general_elements(paths, width, height, colors, style, prompt_lower)
+        
+        return svg_header + '\n' + '\n'.join(paths) + '\n' + svg_footer
+    
+    def _add_icon_elements(self, paths, width, height, colors, style):
+        """Add icon-style elements"""
+        center_x, center_y = width // 2, height // 2
+        
+        if style == "iconography":
+            # Clean geometric icon
+            main_size = min(width, height) // 3
+            paths.append(f'<circle cx="{center_x}" cy="{center_y}" r="{main_size}" fill="none" stroke="{colors[0]}" stroke-width="4"/>')
+            paths.append(f'<circle cx="{center_x}" cy="{center_y}" r="{main_size//2}" fill="{colors[1]}" opacity="0.8"/>')
+        else:
+            # More detailed icon
+            for i in range(3):
+                size = (3-i) * 15
+                paths.append(f'<circle cx="{center_x}" cy="{center_y}" r="{size}" fill="none" stroke="{colors[i]}" stroke-width="3"/>')
+    
+    def _add_landscape_elements(self, paths, width, height, colors, style):
+        """Add landscape elements"""
+        # Sky
+        paths.append(f'<rect x="0" y="0" width="{width}" height="{height//2}" fill="{colors[2]}" opacity="0.3"/>')
+        
+        # Mountains
+        for i in range(3):
+            x1 = i * width // 3
+            x2 = x1 + width // 3
+            peak_x = x1 + width // 6
+            peak_y = height // 3 + i * 10
             
+            points = f"{x1},{height//2} {peak_x},{peak_y} {x2},{height//2}"
+            paths.append(f'<polygon points="{points}" fill="{colors[i % len(colors)]}" opacity="0.7"/>')
+        
+        # Ground
+        paths.append(f'<rect x="0" y="{height//2}" width="{width}" height="{height//2}" fill="{colors[3]}" opacity="0.4"/>')
+    
+    def _add_character_elements(self, paths, width, height, colors, style):
+        """Add character/figure elements"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Head
+        head_r = min(width, height) // 8
+        paths.append(f'<circle cx="{center_x}" cy="{center_y - 30}" r="{head_r}" fill="{colors[0]}" opacity="0.8"/>')
+        
+        # Body
+        body_width = head_r * 2
+        body_height = head_r * 3
+        paths.append(f'<rect x="{center_x - body_width//2}" y="{center_y - 10}" width="{body_width}" height="{body_height}" fill="{colors[1]}" opacity="0.8"/>')
+        
+        # Arms
+        arm_length = head_r * 2
+        paths.append(f'<line x1="{center_x - body_width//2}" y1="{center_y + 10}" x2="{center_x - body_width//2 - arm_length}" y2="{center_y + 30}" stroke="{colors[2]}" stroke-width="4"/>')
+        paths.append(f'<line x1="{center_x + body_width//2}" y1="{center_y + 10}" x2="{center_x + body_width//2 + arm_length}" y2="{center_y + 30}" stroke="{colors[2]}" stroke-width="4"/>')
+    
+    def _add_abstract_elements(self, paths, width, height, colors, style):
+        """Add abstract/geometric elements"""
+        for i in range(8):
+            if i % 3 == 0:
+                # Circles
+                cx = np.random.randint(30, width - 30)
+                cy = np.random.randint(30, height - 30)
+                r = np.random.randint(10, 40)
+                paths.append(f'<circle cx="{cx}" cy="{cy}" r="{r}" fill="{colors[i % len(colors)]}" opacity="0.6"/>')
+            elif i % 3 == 1:
+                # Rectangles
+                x = np.random.randint(10, width - 50)
+                y = np.random.randint(10, height - 50)
+                w = np.random.randint(20, 60)
+                h = np.random.randint(20, 60)
+                paths.append(f'<rect x="{x}" y="{y}" width="{w}" height="{h}" fill="{colors[i % len(colors)]}" opacity="0.6"/>')
             else:
-                # Mixed shapes
-                for i in range(8):
-                    x = np.random.randint(10, width-30)
-                    y = np.random.randint(10, height-30)
-                    color = colors[i % len(colors)]
-                    if i % 2 == 0:
-                        draw.ellipse([x, y, x+20, y+20], fill=color)
-                    else:
-                        draw.rectangle([x, y, x+20, y+20], fill=color)
-            
-            return img
-            
-        except Exception as e:
-            # Return error image
-            img = Image.new('RGB', (224, 224), 'red')
-            return img
+                # Triangles
+                x1 = np.random.randint(20, width - 20)
+                y1 = np.random.randint(20, height - 20)
+                x2 = x1 + np.random.randint(-30, 30)
+                y2 = y1 + np.random.randint(-30, 30)
+                x3 = x1 + np.random.randint(-30, 30)
+                y3 = y1 + np.random.randint(-30, 30)
+                points = f"{x1},{y1} {x2},{y2} {x3},{y3}"
+                paths.append(f'<polygon points="{points}" fill="{colors[i % len(colors)]}" opacity="0.6"/>')
     
+    def _add_general_elements(self, paths, width, height, colors, style, prompt_lower):
+        """Add general elements based on prompt keywords"""
+        # Analyze prompt for specific objects
+        if 'cat' in prompt_lower or 'animal' in prompt_lower:
+            self._add_simple_cat(paths, width, height, colors)
+        elif 'house' in prompt_lower or 'building' in prompt_lower:
+            self._add_simple_house(paths, width, height, colors)
+        elif 'tree' in prompt_lower or 'plant' in prompt_lower:
+            self._add_simple_tree(paths, width, height, colors)
+        elif 'car' in prompt_lower or 'vehicle' in prompt_lower:
+            self._add_simple_car(paths, width, height, colors)
+        else:
+            # Default abstract composition
+            self._add_abstract_elements(paths, width, height, colors, style)
+    
+    def _add_simple_cat(self, paths, width, height, colors):
+        """Add a simple cat figure"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Body
+        paths.append(f'<ellipse cx="{center_x}" cy="{center_y + 20}" rx="40" ry="25" fill="{colors[0]}"/>')
+        # Head
+        paths.append(f'<circle cx="{center_x}" cy="{center_y - 20}" r="25" fill="{colors[0]}"/>')
+        # Ears
+        paths.append(f'<polygon points="{center_x-20},{center_y-35} {center_x-10},{center_y-50} {center_x-5},{center_y-35}" fill="{colors[0]}"/>')
+        paths.append(f'<polygon points="{center_x+5},{center_y-35} {center_x+10},{center_y-50} {center_x+20},{center_y-35}" fill="{colors[0]}"/>')
+        # Tail
+        paths.append(f'<path d="M {center_x+35} {center_y+15} Q {center_x+60} {center_y-10} {center_x+45} {center_y-30}" stroke="{colors[0]}" stroke-width="8" fill="none"/>')
+    
+    def _add_simple_house(self, paths, width, height, colors):
+        """Add a simple house"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Base
+        house_width, house_height = 80, 60
+        paths.append(f'<rect x="{center_x - house_width//2}" y="{center_y}" width="{house_width}" height="{house_height}" fill="{colors[0]}"/>')
+        # Roof
+        roof_points = f"{center_x - house_width//2 - 10},{center_y} {center_x},{center_y - 40} {center_x + house_width//2 + 10},{center_y}"
+        paths.append(f'<polygon points="{roof_points}" fill="{colors[1]}"/>')
+        # Door
+        paths.append(f'<rect x="{center_x - 10}" y="{center_y + 20}" width="20" height="40" fill="{colors[2]}"/>')
+        # Window
+        paths.append(f'<rect x="{center_x + 15}" y="{center_y + 15}" width="15" height="15" fill="{colors[3]}"/>')
+    
+    def _add_simple_tree(self, paths, width, height, colors):
+        """Add a simple tree"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Trunk
+        paths.append(f'<rect x="{center_x - 8}" y="{center_y + 10}" width="16" height="40" fill="{colors[0]}"/>')
+        # Leaves
+        paths.append(f'<circle cx="{center_x}" cy="{center_y - 10}" r="35" fill="{colors[1]}"/>')
+        paths.append(f'<circle cx="{center_x - 20}" cy="{center_y}" r="25" fill="{colors[1]}"/>')
+        paths.append(f'<circle cx="{center_x + 20}" cy="{center_y}" r="25" fill="{colors[1]}"/>')
+    
+    def _add_simple_car(self, paths, width, height, colors):
+        """Add a simple car"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Body
+        paths.append(f'<rect x="{center_x - 50}" y="{center_y}" width="100" height="30" fill="{colors[0]}"/>')
+        # Top
+        paths.append(f'<rect x="{center_x - 30}" y="{center_y - 20}" width="60" height="20" fill="{colors[0]}"/>')
+        # Wheels
+        paths.append(f'<circle cx="{center_x - 30}" cy="{center_y + 35}" r="12" fill="{colors[1]}"/>')
+        paths.append(f'<circle cx="{center_x + 30}" cy="{center_y + 35}" r="12" fill="{colors[1]}"/>')
+        # Windows
+        paths.append(f'<rect x="{center_x - 25}" y="{center_y - 15}" width="20" height="12" fill="{colors[2]}"/>')
+        paths.append(f'<rect x="{center_x + 5}" y="{center_y - 15}" width="20" height="12" fill="{colors[2]}"/>')
+
     def _generate_simple_svg(self, prompt: str, width: int, height: int, particle_id: int, style: str) -> str:
         """
         Generate a simple SVG as placeholder for each particle