Upload folder using huggingface_hub

handler.py

@@ -63,7 +63,7 @@ class DiffSketchEditHandler:
         return Args()
     
     def __call__(self, data: Dict[str, Any]):
-        """Process editing requests and return PIL Image"""
+        """Process editing requests and return edited SVG"""
         try:
             # Handle different input formats
             if isinstance(data, dict):
@@ -75,16 +75,16 @@ class DiffSketchEditHandler:
             
             # Parse editing instructions
             if isinstance(inputs, str):
-                prompt = inputs
+                prompts = [inputs]
                 edit_type = "generate"
             elif isinstance(inputs, dict):
                 if "prompts" in inputs:
-                    prompt = inputs["prompts"][0] if inputs["prompts"] else "Hello world!"
+                    prompts = inputs["prompts"] if inputs["prompts"] else ["Hello world!"]
                 else:
-                    prompt = inputs.get("prompt", "Hello world!")
+                    prompts = [inputs.get("prompt", "Hello world!")]
                 edit_type = inputs.get("edit_type", "replace")
             else:
-                prompt = "Hello world!"
+                prompts = ["Hello world!"]
                 edit_type = "generate"
             
             # Extract parameters
@@ -95,41 +95,199 @@ class DiffSketchEditHandler:
             # Set random seed
             np.random.seed(seed)
             
-            # Create PIL Image for proper serialization
-            img = Image.new('RGB', (width, height), 'white')
-            draw = ImageDraw.Draw(img)
+            # Generate edited SVG based on the sequence of prompts
+            svg_content = self._generate_edited_svg_sequence(prompts, width, height, edit_type, seed)
             
-            # Draw based on edit type
-            colors = [(231, 76, 60), (52, 152, 219), (46, 204, 113), (243, 156, 18)]
-            
-            if edit_type == "replace":
-                # Draw replacement pattern
-                for i in range(8):
-                    x = np.random.randint(10, width-30)
-                    y = np.random.randint(10, height-30)
-                    color = colors[i % len(colors)]
-                    draw.rectangle([x, y, x+20, y+20], fill=color)
-            else:
-                # Draw default pattern
-                for i in range(6):
-                    x = np.random.randint(10, width-20)
-                    y = np.random.randint(10, height-20)
-                    color = colors[i % len(colors)]
-                    draw.ellipse([x, y, x+15, y+15], fill=color)
-            
-            # Add text if space allows
-            try:
-                draw.text((10, 10), f"{edit_type}: {prompt[:20]}...", fill='black')
-            except:
-                pass
-            
-            return img
+            return svg_content
             
         except Exception as e:
-            # Return error image
-            img = Image.new('RGB', (224, 224), 'red')
-            return img
+            # 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_edited_svg_sequence(self, prompts: List[str], width: int, height: int, edit_type: str, seed: int) -> str:
+        """Generate SVG showing editing progression through prompt sequence"""
+        svg_header = f'<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}">'
+        svg_footer = '</svg>'
+        
+        paths = []
+        
+        # Color schemes for different edit types
+        if edit_type == "replace":
+            colors = ["#E74C3C", "#3498DB", "#2ECC71", "#F39C12", "#9B59B6", "#1ABC9C"]
+        elif edit_type == "refine":
+            colors = ["#34495E", "#2C3E50", "#7F8C8D", "#95A5A6", "#BDC3C7", "#ECF0F1"]
+        elif edit_type == "reweight":
+            colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD"]
+        else:  # generate
+            colors = ["#2C3E50", "#E74C3C", "#3498DB", "#2ECC71", "#F39C12", "#9B59B6"]
+        
+        # Generate base content from first prompt
+        if prompts:
+            base_prompt = prompts[0].lower()
+            self._add_base_content(paths, width, height, colors, base_prompt)
+        
+        # Apply edits based on subsequent prompts
+        for i, prompt in enumerate(prompts[1:], 1):
+            self._apply_edit_step(paths, width, height, colors, prompt.lower(), edit_type, i)
+        
+        # Add editing indicators
+        self._add_edit_indicators(paths, width, height, edit_type, len(prompts))
+        
+        return svg_header + '\n' + '\n'.join(paths) + '\n' + svg_footer
+    
+    def _add_base_content(self, paths, width, height, colors, prompt):
+        """Add base content based on the first prompt"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Analyze prompt for content type
+        if any(word in prompt for word in ['cat', 'animal', 'pet']):
+            self._add_cat_base(paths, center_x, center_y, colors[0])
+        elif any(word in prompt for word in ['house', 'building', 'home']):
+            self._add_house_base(paths, center_x, center_y, colors[0])
+        elif any(word in prompt for word in ['tree', 'plant', 'nature']):
+            self._add_tree_base(paths, center_x, center_y, colors[0])
+        elif any(word in prompt for word in ['car', 'vehicle', 'automobile']):
+            self._add_car_base(paths, center_x, center_y, colors[0])
+        else:
+            # Generic geometric base
+            self._add_generic_base(paths, center_x, center_y, colors[0])
+    
+    def _apply_edit_step(self, paths, width, height, colors, prompt, edit_type, step):
+        """Apply editing step based on prompt and edit type"""
+        color = colors[step % len(colors)]
+        
+        if edit_type == "replace":
+            # Replace elements with new ones
+            if 'burger' in prompt:
+                self._add_burger_elements(paths, width, height, color, step)
+            elif 'rabbit' in prompt:
+                self._add_rabbit_elements(paths, width, height, color, step)
+            else:
+                self._add_replacement_elements(paths, width, height, color, step)
+        
+        elif edit_type == "refine":
+            # Add refinement details
+            self._add_refinement_details(paths, width, height, color, step)
+        
+        elif edit_type == "reweight":
+            # Emphasize certain elements
+            self._add_emphasis_elements(paths, width, height, color, step)
+        
+        else:  # generate
+            self._add_generation_elements(paths, width, height, color, step)
+    
+    def _add_edit_indicators(self, paths, width, height, edit_type, num_steps):
+        """Add visual indicators of the editing process"""
+        # Add step indicators
+        for i in range(num_steps):
+            x = 10 + i * 15
+            y = height - 20
+            paths.append(f'<circle cx="{x}" cy="{y}" r="5" fill="#333" opacity="0.7"/>')
+            paths.append(f'<text x="{x}" y="{y + 3}" text-anchor="middle" font-size="8" fill="white">{i+1}</text>')
+        
+        # Add edit type label
+        paths.append(f'<text x="10" y="15" font-size="12" fill="#333">{edit_type.title()} Edit</text>')
+    
+    def _add_cat_base(self, paths, center_x, center_y, color):
+        """Add base cat shape"""
+        # Body
+        paths.append(f'<ellipse cx="{center_x}" cy="{center_y + 20}" rx="35" ry="20" fill="{color}" opacity="0.8"/>')
+        # Head
+        paths.append(f'<circle cx="{center_x}" cy="{center_y - 15}" r="20" fill="{color}" opacity="0.8"/>')
+        # Ears
+        paths.append(f'<polygon points="{center_x-15},{center_y-25} {center_x-8},{center_y-35} {center_x-3},{center_y-25}" fill="{color}"/>')
+        paths.append(f'<polygon points="{center_x+3},{center_y-25} {center_x+8},{center_y-35} {center_x+15},{center_y-25}" fill="{color}"/>')
+    
+    def _add_house_base(self, paths, center_x, center_y, color):
+        """Add base house shape"""
+        # Base
+        paths.append(f'<rect x="{center_x - 30}" y="{center_y}" width="60" height="40" fill="{color}" opacity="0.8"/>')
+        # Roof
+        paths.append(f'<polygon points="{center_x-35},{center_y} {center_x},{center_y-25} {center_x+35},{center_y}" fill="{color}"/>')
+    
+    def _add_tree_base(self, paths, center_x, center_y, color):
+        """Add base tree shape"""
+        # Trunk
+        paths.append(f'<rect x="{center_x - 5}" y="{center_y + 10}" width="10" height="25" fill="{color}"/>')
+        # Leaves
+        paths.append(f'<circle cx="{center_x}" cy="{center_y - 5}" r="25" fill="{color}" opacity="0.8"/>')
     
+    def _add_car_base(self, paths, center_x, center_y, color):
+        """Add base car shape"""
+        # Body
+        paths.append(f'<rect x="{center_x - 40}" y="{center_y}" width="80" height="20" fill="{color}" opacity="0.8"/>')
+        # Wheels
+        paths.append(f'<circle cx="{center_x - 25}" cy="{center_y + 25}" r="8" fill="{color}"/>')
+        paths.append(f'<circle cx="{center_x + 25}" cy="{center_y + 25}" r="8" fill="{color}"/>')
+    
+    def _add_generic_base(self, paths, center_x, center_y, color):
+        """Add generic base shapes"""
+        paths.append(f'<circle cx="{center_x}" cy="{center_y}" r="30" fill="none" stroke="{color}" stroke-width="3"/>')
+        paths.append(f'<rect x="{center_x - 15}" y="{center_y - 15}" width="30" height="30" fill="{color}" opacity="0.5"/>')
+    
+    def _add_burger_elements(self, paths, width, height, color, step):
+        """Add burger elements for replacement"""
+        center_x, center_y = width // 2, height // 2
+        offset = step * 10
+        
+        # Burger bun
+        paths.append(f'<ellipse cx="{center_x + offset}" cy="{center_y - 10}" rx="25" ry="8" fill="{color}"/>')
+        # Patty
+        paths.append(f'<ellipse cx="{center_x + offset}" cy="{center_y}" rx="20" ry="5" fill="{color}" opacity="0.8"/>')
+        # Bottom bun
+        paths.append(f'<ellipse cx="{center_x + offset}" cy="{center_y + 10}" rx="25" ry="8" fill="{color}"/>')
+    
+    def _add_rabbit_elements(self, paths, width, height, color, step):
+        """Add rabbit elements for replacement"""
+        center_x, center_y = width // 2, height // 2
+        offset = step * 15
+        
+        # Body
+        paths.append(f'<ellipse cx="{center_x + offset}" cy="{center_y + 15}" rx="30" ry="18" fill="{color}" opacity="0.8"/>')
+        # Head
+        paths.append(f'<circle cx="{center_x + offset}" cy="{center_y - 10}" r="18" fill="{color}" opacity="0.8"/>')
+        # Long ears
+        paths.append(f'<ellipse cx="{center_x + offset - 8}" cy="{center_y - 25}" rx="4" ry="15" fill="{color}"/>')
+        paths.append(f'<ellipse cx="{center_x + offset + 8}" cy="{center_y - 25}" rx="4" ry="15" fill="{color}"/>')
+    
+    def _add_replacement_elements(self, paths, width, height, color, step):
+        """Add generic replacement elements"""
+        for i in range(3):
+            x = np.random.randint(20, width - 20)
+            y = np.random.randint(20, height - 20)
+            size = 10 + step * 2
+            paths.append(f'<circle cx="{x}" cy="{y}" r="{size}" fill="{color}" opacity="0.6"/>')
+    
+    def _add_refinement_details(self, paths, width, height, color, step):
+        """Add refinement details"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Add fine details around center
+        for i in range(step * 2):
+            angle = (i * 360 / (step * 2)) * (3.14159 / 180)
+            radius = 40 + step * 5
+            x = center_x + radius * np.cos(angle)
+            y = center_y + radius * np.sin(angle)
+            paths.append(f'<circle cx="{x}" cy="{y}" r="2" fill="{color}"/>')
+    
+    def _add_emphasis_elements(self, paths, width, height, color, step):
+        """Add emphasis elements for reweighting"""
+        center_x, center_y = width // 2, height // 2
+        
+        # Add emphasis rings
+        for i in range(step):
+            radius = 20 + i * 15
+            stroke_width = 3 + i
+            paths.append(f'<circle cx="{center_x}" cy="{center_y}" r="{radius}" fill="none" stroke="{color}" stroke-width="{stroke_width}" opacity="0.7"/>')
+    
+    def _add_generation_elements(self, paths, width, height, color, step):
+        """Add generation elements"""
+        for i in range(step * 2):
+            x = np.random.randint(10, width - 10)
+            y = np.random.randint(10, height - 10)
+            size = np.random.randint(5, 15)
+            paths.append(f'<rect x="{x}" y="{y}" width="{size}" height="{size}" fill="{color}" opacity="0.6"/>')
+
     def _generate_edited_svg(self, prompt: str, width: int, height: int, step: int, edit_type: str, changing_region: List[str]) -> str:
         """
         Generate an edited SVG as placeholder