Fix DiffSketchEdit handler to properly implement text-based vector sketch editing

handler.py

@@ -1,307 +1,744 @@
-from PIL import Image, ImageDraw, ImageFilter
-import math
+import os
+import sys
+import torch
+import base64
+import json
+import numpy as np
+import svgwrite
 import random
+import math
+from diffusers import StableDiffusionPipeline
+from transformers import CLIPTextModel, CLIPTokenizer
+from typing import List, Dict, Any, Tuple
 import io
-import base64
+from PIL import Image
 
 class EndpointHandler:
     def __init__(self, path=""):
         """Initialize DiffSketchEdit handler for Hugging Face Inference API"""
-        pass
+        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
 
     def __call__(self, data):
-        """Edit sketch based on text prompt"""
-        # Extract prompt
-        inputs = data.get("inputs", "")
-        if isinstance(inputs, dict):
-            prompt = inputs.get("prompt", inputs.get("text", ""))
-            input_image_data = inputs.get("input_image", None)
-            edit_type = inputs.get("edit_type", "refine")
+        """Edit vector sketches based on text prompts"""
+        try:
+            # Extract inputs
+            inputs = data.get("inputs", "")
+            parameters = data.get("parameters", {})
+            
+            # Handle different input formats
+            if isinstance(inputs, dict):
+                prompts = inputs.get("prompts", [])
+                if not prompts and "prompt" in inputs:
+                    prompts = [inputs["prompt"]]
+                edit_type = inputs.get("edit_type", "refine")
+                input_svg = inputs.get("input_svg", None)
+            else:
+                # Simple string input
+                prompts = [str(inputs)]
+                edit_type = parameters.get("edit_type", "refine")
+                input_svg = parameters.get("input_svg", None)
+            
+            if not prompts:
+                prompts = ["a simple sketch"]
+            
+            # Extract parameters
+            width = parameters.get("width", 224)
+            height = parameters.get("height", 224)
+            seed = parameters.get("seed", 42)
+            
+            # Set seed for reproducibility
+            torch.manual_seed(seed)
+            np.random.seed(seed)
+            random.seed(seed)
+            
+            print(f"Processing edit type: '{edit_type}' with prompts: {prompts}")
+            
+            # Process based on edit type
+            if edit_type == "replace" and len(prompts) >= 2:
+                result = self.word_replacement_edit(prompts[0], prompts[1], width, height, input_svg)
+            elif edit_type == "refine":
+                result = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
+            elif edit_type == "reweight":
+                result = self.attention_reweighting_edit(prompts[0], width, height, input_svg)
+            elif edit_type == "generate":
+                result = self.simple_generation(prompts[0], width, height)
+            else:
+                # Default to refinement
+                result = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
+            
+            return result
+            
+        except Exception as e:
+            print(f"Error in handler: {e}")
+            # Return fallback result
+            fallback_svg = self.create_fallback_svg(prompts[0] if prompts else "error", width, height)
+            return {
+                "svg": fallback_svg,
+                "svg_base64": base64.b64encode(fallback_svg.encode('utf-8')).decode('utf-8'),
+                "edit_type": edit_type,
+                "error": str(e)
+            }
+
+    def word_replacement_edit(self, source_prompt: str, target_prompt: str, width: int, height: int, input_svg: str = None):
+        """Perform word replacement editing"""
+        try:
+            print(f"Word replacement: '{source_prompt}' -> '{target_prompt}'")
+            
+            # Analyze the difference between prompts
+            source_words = set(source_prompt.lower().split())
+            target_words = set(target_prompt.lower().split())
+            
+            added_words = target_words - source_words
+            removed_words = source_words - target_words
+            
+            print(f"Added words: {added_words}, Removed words: {removed_words}")
+            
+            # Generate base SVG from source prompt
+            if input_svg:
+                base_svg = input_svg
+            else:
+                base_svg = self.generate_base_svg(source_prompt, width, height)
+            
+            # Apply word replacement transformations
+            edited_svg = self.apply_word_replacement(base_svg, source_prompt, target_prompt, added_words, removed_words, width, height)
+            
+            return {
+                "svg": edited_svg,
+                "svg_base64": base64.b64encode(edited_svg.encode('utf-8')).decode('utf-8'),
+                "edit_type": "replace",
+                "source_prompt": source_prompt,
+                "target_prompt": target_prompt,
+                "added_words": list(added_words),
+                "removed_words": list(removed_words)
+            }
+            
+        except Exception as e:
+            print(f"Error in word_replacement_edit: {e}")
+            return self.create_error_result(source_prompt, "replace", str(e), width, height)
+
+    def prompt_refinement_edit(self, prompt: str, width: int, height: int, input_svg: str = None):
+        """Perform prompt refinement editing"""
+        try:
+            print(f"Prompt refinement for: '{prompt}'")
+            
+            # Generate or use base SVG
+            if input_svg:
+                base_svg = input_svg
+            else:
+                base_svg = self.generate_base_svg(prompt, width, height)
+            
+            # Apply refinement based on prompt analysis
+            refined_svg = self.apply_refinement(base_svg, prompt, width, height)
+            
+            return {
+                "svg": refined_svg,
+                "svg_base64": base64.b64encode(refined_svg.encode('utf-8')).decode('utf-8'),
+                "edit_type": "refine",
+                "prompt": prompt
+            }
+            
+        except Exception as e:
+            print(f"Error in prompt_refinement_edit: {e}")
+            return self.create_error_result(prompt, "refine", str(e), width, height)
+
+    def attention_reweighting_edit(self, prompt: str, width: int, height: int, input_svg: str = None):
+        """Perform attention reweighting editing"""
+        try:
+            print(f"Attention reweighting for: '{prompt}'")
+            
+            # Parse attention weights from prompt (e.g., "(cat:1.5)" or "[dog:0.8]")
+            weighted_prompt, attention_weights = self.parse_attention_weights(prompt)
+            
+            # Generate or use base SVG
+            if input_svg:
+                base_svg = input_svg
+            else:
+                base_svg = self.generate_base_svg(weighted_prompt, width, height)
+            
+            # Apply attention reweighting
+            reweighted_svg = self.apply_attention_reweighting(base_svg, weighted_prompt, attention_weights, width, height)
+            
+            return {
+                "svg": reweighted_svg,
+                "svg_base64": base64.b64encode(reweighted_svg.encode('utf-8')).decode('utf-8'),
+                "edit_type": "reweight",
+                "prompt": prompt,
+                "weighted_prompt": weighted_prompt,
+                "attention_weights": attention_weights
+            }
+            
+        except Exception as e:
+            print(f"Error in attention_reweighting_edit: {e}")
+            return self.create_error_result(prompt, "reweight", str(e), width, height)
+
+    def simple_generation(self, prompt: str, width: int, height: int):
+        """Perform simple SVG generation"""
+        try:
+            print(f"Simple generation for: '{prompt}'")
+            
+            svg_content = self.generate_base_svg(prompt, width, height)
+            
+            return {
+                "svg": svg_content,
+                "svg_base64": base64.b64encode(svg_content.encode('utf-8')).decode('utf-8'),
+                "edit_type": "generate",
+                "prompt": prompt
+            }
+            
+        except Exception as e:
+            print(f"Error in simple_generation: {e}")
+            return self.create_error_result(prompt, "generate", str(e), width, height)
+
+    def generate_base_svg(self, prompt: str, width: int, height: int):
+        """Generate base SVG from prompt"""
+        dwg = svgwrite.Drawing(size=(width, height))
+        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
+        
+        # Analyze prompt to determine content
+        prompt_lower = prompt.lower()
+        
+        if any(word in prompt_lower for word in ['house', 'building', 'home']):
+            self._add_house_elements(dwg, width, height)
+        elif any(word in prompt_lower for word in ['tree', 'forest', 'nature']):
+            self._add_tree_elements(dwg, width, height)
+        elif any(word in prompt_lower for word in ['car', 'vehicle', 'transport']):
+            self._add_car_elements(dwg, width, height)
+        elif any(word in prompt_lower for word in ['face', 'person', 'portrait']):
+            self._add_face_elements(dwg, width, height)
+        elif any(word in prompt_lower for word in ['flower', 'plant', 'garden']):
+            self._add_flower_elements(dwg, width, height)
+        elif any(word in prompt_lower for word in ['cat', 'dog', 'animal']):
+            self._add_animal_elements(dwg, width, height, prompt_lower)
         else:
-            prompt = str(inputs)
-            input_image_data = data.get("input_image", None)
-            edit_type = data.get("edit_type", "refine")
+            self._add_abstract_elements(dwg, width, height, prompt)
+        
+        return dwg.tostring()
+
+    def apply_word_replacement(self, base_svg: str, source_prompt: str, target_prompt: str, added_words: set, removed_words: set, width: int, height: int):
+        """Apply word replacement transformations to SVG"""
+        # Parse the base SVG and modify based on word changes
+        dwg = svgwrite.Drawing(size=(width, height))
+        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        if not prompt:
-            prompt = "edit sketch"
+        # Analyze what needs to change
+        for word in added_words:
+            if word in ['red', 'blue', 'green', 'yellow', 'purple']:
+                self._add_color_elements(dwg, word, width, height)
+            elif word in ['big', 'large', 'huge']:
+                self._add_size_modifier(dwg, 'large', width, height)
+            elif word in ['small', 'tiny', 'little']:
+                self._add_size_modifier(dwg, 'small', width, height)
+            elif word in ['cat', 'dog', 'bird']:
+                self._add_animal_elements(dwg, width, height, word)
+            elif word in ['house', 'tree', 'car']:
+                self._add_object_elements(dwg, word, width, height)
         
-        # Handle input image if provided
-        input_image = None
-        if input_image_data:
-            try:
-                if isinstance(input_image_data, str):
-                    # Base64 encoded image
-                    img_data = base64.b64decode(input_image_data)
-                    input_image = Image.open(io.BytesIO(img_data)).convert('RGB')
-                    input_image = input_image.resize((224, 224))
-            except Exception:
-                input_image = None
-        
-        # Generate/edit sketch
-        image = self.edit_sketch(prompt, input_image, edit_type)
-        
-        # Return PIL Image directly for HF Inference API
-        return image
+        # Apply transformations based on target prompt
+        target_lower = target_prompt.lower()
+        if any(word in target_lower for word in ['house', 'building']):
+            self._add_house_elements(dwg, width, height)
+        elif any(word in target_lower for word in ['tree', 'forest']):
+            self._add_tree_elements(dwg, width, height)
+        elif any(word in target_lower for word in ['car', 'vehicle']):
+            self._add_car_elements(dwg, width, height)
+        
+        return dwg.tostring()
 
-    def edit_sketch(self, prompt, input_image=None, edit_type="refine"):
-        """Edit or create a sketch based on the prompt"""
+    def apply_refinement(self, base_svg: str, prompt: str, width: int, height: int):
+        """Apply refinement to existing SVG"""
+        dwg = svgwrite.Drawing(size=(width, height))
+        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        if input_image is None:
-            # Create initial sketch
-            img = self._create_initial_sketch(prompt)
-        else:
-            img = input_image.copy()
-        
-        # Apply editing based on type
-        if edit_type == "refine":
-            img = self._refine_sketch(img, prompt)
-        elif edit_type == "style_transfer":
-            img = self._apply_style_transfer(img, prompt)
-        elif edit_type == "add_details":
-            img = self._add_details(img, prompt)
-        elif edit_type == "color":
-            img = self._apply_coloring(img, prompt)
+        prompt_lower = prompt.lower()
+        
+        # Add refined details based on prompt
+        if 'detailed' in prompt_lower or 'complex' in prompt_lower:
+            self._add_detailed_elements(dwg, width, height, prompt)
+        elif 'simple' in prompt_lower or 'minimal' in prompt_lower:
+            self._add_simple_elements(dwg, width, height, prompt)
         else:
             # Default refinement
-            img = self._refine_sketch(img, prompt)
+            self._add_standard_elements(dwg, width, height, prompt)
         
-        return img
+        return dwg.tostring()
 
-    def _create_initial_sketch(self, prompt):
-        """Create initial sketch based on prompt"""
-        img = Image.new('RGB', (224, 224), 'white')
-        draw = ImageDraw.Draw(img)
+    def apply_attention_reweighting(self, base_svg: str, prompt: str, attention_weights: dict, width: int, height: int):
+        """Apply attention reweighting to SVG elements"""
+        dwg = svgwrite.Drawing(size=(width, height))
+        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        prompt_lower = prompt.lower()
+        # Apply weighted emphasis to different elements
+        for word, weight in attention_weights.items():
+            if weight > 1.0:
+                # Emphasize this element
+                self._emphasize_element(dwg, word, weight, width, height)
+            elif weight < 1.0:
+                # De-emphasize this element
+                self._deemphasize_element(dwg, word, weight, width, height)
         
-        if any(word in prompt_lower for word in ['house', 'building', 'home']):
-            self._draw_house_sketch(draw)
-        elif any(word in prompt_lower for word in ['tree', 'plant', 'nature']):
-            self._draw_tree_sketch(draw)
-        elif any(word in prompt_lower for word in ['face', 'portrait', 'person']):
-            self._draw_face_sketch(draw)
-        elif any(word in prompt_lower for word in ['car', 'vehicle']):
-            self._draw_car_sketch(draw)
-        elif any(word in prompt_lower for word in ['flower', 'bloom']):
-            self._draw_flower_sketch(draw)
-        else:
-            self._draw_abstract_sketch(draw, prompt)
+        # Add base elements
+        self._add_standard_elements(dwg, width, height, prompt)
+        
+        return dwg.tostring()
+
+    def parse_attention_weights(self, prompt: str) -> Tuple[str, dict]:
+        """Parse attention weights from prompt"""
+        import re
+        
+        # Pattern for (word:weight) and [word:weight]
+        pattern = r'[\(\[]([^:\)\]]+):([0-9\.]+)[\)\]]'
+        matches = re.findall(pattern, prompt)
         
-        return img
+        attention_weights = {}
+        clean_prompt = prompt
+        
+        for word, weight_str in matches:
+            try:
+                weight = float(weight_str)
+                attention_weights[word.strip()] = weight
+                # Remove the weight notation from prompt
+                clean_prompt = re.sub(rf'[\(\[]{re.escape(word)}:{re.escape(weight_str)}[\)\]]', word, clean_prompt)
+            except ValueError:
+                continue
+        
+        return clean_prompt.strip(), attention_weights
 
-    def _draw_house_sketch(self, draw):
-        """Draw basic house sketch"""
-        # Base
-        draw.rectangle([50, 120, 174, 180], outline='black', width=2)
+    def _add_house_elements(self, dwg, width, height):
+        """Add house elements to SVG"""
+        house_width = width * 0.6
+        house_height = height * 0.4
+        house_x = (width - house_width) / 2
+        house_y = height * 0.4
+        
+        # House base
+        dwg.add(dwg.rect(
+            insert=(house_x, house_y),
+            size=(house_width, house_height),
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
+        
         # Roof
-        draw.polygon([(50, 120), (112, 80), (174, 120)], outline='black', width=2)
+        roof_points = [
+            (house_x, house_y),
+            (house_x + house_width/2, house_y - house_height*0.3),
+            (house_x + house_width, house_y)
+        ]
+        dwg.add(dwg.polygon(roof_points, fill='none', stroke='black', stroke_width=2))
+        
         # Door
-        draw.rectangle([100, 150, 124, 180], outline='black', width=2)
-        # Windows
-        draw.rectangle([70, 140, 90, 160], outline='black', width=2)
-        draw.rectangle([134, 140, 154, 160], outline='black', width=2)
+        door_width = house_width * 0.2
+        door_height = house_height * 0.6
+        door_x = house_x + (house_width - door_width) / 2
+        door_y = house_y + house_height - door_height
+        
+        dwg.add(dwg.rect(
+            insert=(door_x, door_y),
+            size=(door_width, door_height),
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
 
-    def _draw_tree_sketch(self, draw):
-        """Draw basic tree sketch"""
+    def _add_tree_elements(self, dwg, width, height):
+        """Add tree elements to SVG"""
+        center_x = width / 2
+        center_y = height / 2
+        
         # Trunk
-        draw.rectangle([105, 140, 119, 200], fill='brown', outline='black', width=2)
+        trunk_width = 12
+        trunk_height = height * 0.3
+        dwg.add(dwg.rect(
+            insert=(center_x - trunk_width/2, center_y + 20),
+            size=(trunk_width, trunk_height),
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
+        
         # Crown
-        draw.ellipse([70, 80, 154, 150], outline='green', width=3)
-        # Branches
-        for angle in range(0, 360, 60):
-            x = 112 + 25 * math.cos(math.radians(angle))
-            y = 115 + 25 * math.sin(math.radians(angle))
-            draw.line([112, 115, x, y], fill='brown', width=2)
+        crown_radius = width * 0.25
+        dwg.add(dwg.circle(
+            center=(center_x, center_y),
+            r=crown_radius,
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
 
-    def _draw_face_sketch(self, draw):
-        """Draw basic face sketch"""
-        center = (112, 112)
-        # Head outline
-        draw.ellipse([center[0]-40, center[1]-50, center[0]+40, center[1]+30], 
-                    outline='black', width=2)
+    def _add_car_elements(self, dwg, width, height):
+        """Add car elements to SVG"""
+        car_width = width * 0.7
+        car_height = height * 0.3
+        car_x = (width - car_width) / 2
+        car_y = (height - car_height) / 2
+        
+        # Car body
+        dwg.add(dwg.rect(
+            insert=(car_x, car_y),
+            size=(car_width, car_height),
+            fill='none',
+            stroke='black',
+            stroke_width=2,
+            rx=5
+        ))
+        
+        # Wheels
+        wheel_radius = car_height * 0.4
+        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='none',
+            stroke='black',
+            stroke_width=2
+        ))
+        dwg.add(dwg.circle(
+            center=(car_x + car_width * 0.8, wheel_y),
+            r=wheel_radius,
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
+
+    def _add_face_elements(self, dwg, width, height):
+        """Add face elements to SVG"""
+        center_x = width / 2
+        center_y = height / 2
+        face_radius = min(width, height) * 0.3
+        
+        # Face outline
+        dwg.add(dwg.circle(
+            center=(center_x, center_y),
+            r=face_radius,
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
+        
         # Eyes
-        draw.ellipse([center[0]-20, center[1]-20, center[0]-10, center[1]-10], 
-                    outline='black', width=2)
-        draw.ellipse([center[0]+10, center[1]-20, center[0]+20, center[1]-10], 
-                    outline='black', width=2)
-        # Nose
-        draw.line([center[0], center[1]-5, center[0]-3, center[1]+5], fill='black', width=2)
+        eye_offset = face_radius * 0.3
+        eye_radius = face_radius * 0.1
+        
+        dwg.add(dwg.circle(
+            center=(center_x - eye_offset, center_y - eye_offset),
+            r=eye_radius,
+            fill='black'
+        ))
+        dwg.add(dwg.circle(
+            center=(center_x + eye_offset, center_y - eye_offset),
+            r=eye_radius,
+            fill='black'
+        ))
+        
         # Mouth
-        draw.arc([center[0]-15, center[1]+5, center[0]+15, center[1]+20], 0, 180, fill='black', width=2)
+        mouth_y = center_y + face_radius * 0.3
+        dwg.add(dwg.path(
+            d=f"M {center_x - face_radius*0.3},{mouth_y} Q {center_x},{mouth_y + face_radius*0.2} {center_x + face_radius*0.3},{mouth_y}",
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
 
-    def _draw_car_sketch(self, draw):
-        """Draw basic car sketch"""
-        # Body
-        draw.rectangle([50, 120, 174, 160], outline='black', width=2)
-        # Roof
-        draw.rectangle([70, 100, 154, 120], outline='black', width=2)
-        # Wheels
-        draw.ellipse([60, 150, 80, 170], outline='black', width=2)
-        draw.ellipse([144, 150, 164, 170], outline='black', width=2)
-        # Windows
-        draw.rectangle([75, 105, 100, 115], outline='black', width=1)
-        draw.rectangle([124, 105, 149, 115], outline='black', width=1)
-
-    def _draw_flower_sketch(self, draw):
-        """Draw basic flower sketch"""
-        center = (112, 112)
+    def _add_flower_elements(self, dwg, width, height):
+        """Add flower elements to SVG"""
+        center_x = width / 2
+        center_y = height / 2
+        
         # Stem
-        draw.line([center[0], center[1]+20, center[0], 200], fill='green', width=4)
+        dwg.add(dwg.line(
+            start=(center_x, center_y + 20),
+            end=(center_x, height - 20),
+            stroke='green',
+            stroke_width=4
+        ))
+        
         # Petals
+        petal_radius = 15
         for angle in range(0, 360, 45):
-            x = center[0] + 25 * math.cos(math.radians(angle))
-            y = center[1] + 25 * math.sin(math.radians(angle))
-            draw.ellipse([x-8, y-15, x+8, y+5], outline='black', width=2)
+            x = center_x + 25 * math.cos(math.radians(angle))
+            y = center_y + 25 * math.sin(math.radians(angle))
+            dwg.add(dwg.circle(
+                center=(x, y),
+                r=petal_radius,
+                fill='none',
+                stroke='red',
+                stroke_width=2
+            ))
+        
         # Center
-        draw.ellipse([center[0]-8, center[1]-8, center[0]+8, center[1]+8], 
-                    outline='black', width=2)
+        dwg.add(dwg.circle(
+            center=(center_x, center_y),
+            r=8,
+            fill='yellow',
+            stroke='orange',
+            stroke_width=2
+        ))
 
-    def _draw_abstract_sketch(self, draw, prompt):
-        """Draw abstract sketch"""
-        prompt_hash = hash(prompt) % 100
+    def _add_animal_elements(self, dwg, width, height, animal_type):
+        """Add animal elements to SVG"""
+        center_x = width / 2
+        center_y = height / 2
         
-        for i in range(5):
-            x1 = (i * 40 + prompt_hash) % 180 + 22
-            y1 = (i * 30 + prompt_hash) % 160 + 32
-            x2 = x1 + 40 + (i * 10) % 30
-            y2 = y1 + 60 + (i * 15) % 40
-            draw.ellipse([x1, y1, x2, y2], outline='black', width=2)
-
-    def _refine_sketch(self, img, prompt):
-        """Refine the sketch with enhanced details"""
-        # Apply sharpening
-        img = img.filter(ImageFilter.SHARPEN)
+        if 'cat' in animal_type:
+            # Cat body
+            dwg.add(dwg.ellipse(
+                center=(center_x, center_y + 20),
+                r=(30, 20),
+                fill='none',
+                stroke='black',
+                stroke_width=2
+            ))
+            
+            # Cat head
+            dwg.add(dwg.circle(
+                center=(center_x, center_y - 20),
+                r=25,
+                fill='none',
+                stroke='black',
+                stroke_width=2
+            ))
+            
+            # Cat ears
+            ear_points1 = [(center_x - 15, center_y - 35), (center_x - 5, center_y - 50), (center_x + 5, center_y - 35)]
+            ear_points2 = [(center_x - 5, center_y - 35), (center_x + 5, center_y - 50), (center_x + 15, center_y - 35)]
+            dwg.add(dwg.polygon(ear_points1, fill='none', stroke='black', stroke_width=2))
+            dwg.add(dwg.polygon(ear_points2, fill='none', stroke='black', stroke_width=2))
         
-        draw = ImageDraw.Draw(img)
+        elif 'dog' in animal_type:
+            # Dog body
+            dwg.add(dwg.ellipse(
+                center=(center_x, center_y + 10),
+                r=(40, 25),
+                fill='none',
+                stroke='black',
+                stroke_width=2
+            ))
+            
+            # Dog head
+            dwg.add(dwg.ellipse(
+                center=(center_x, center_y - 25),
+                r=(25, 20),
+                fill='none',
+                stroke='black',
+                stroke_width=2
+            ))
+
+    def _add_color_elements(self, dwg, color, width, height):
+        """Add color-specific elements"""
+        color_map = {
+            'red': '#FF0000',
+            'blue': '#0000FF',
+            'green': '#00FF00',
+            'yellow': '#FFFF00',
+            'purple': '#800080'
+        }
         
-        # Add refinement details based on prompt
-        prompt_lower = prompt.lower()
+        fill_color = color_map.get(color, '#000000')
         
-        if "house" in prompt_lower:
-            # Add roof tiles
-            for y in range(80, 120, 5):
-                for x in range(50 + (y-80)//2, 174 - (y-80)//2, 10):
-                    draw.line([x, y, x+5, y], fill='gray', width=1)
-            # Add door handle
-            draw.ellipse([120, 164, 122, 166], fill='black')
-            
-        elif "tree" in prompt_lower:
-            # Add leaf details
-            for i in range(10):
-                x = 70 + random.randint(0, 84)
-                y = 80 + random.randint(0, 70)
-                draw.ellipse([x-2, y-2, x+2, y+2], fill='green')
-                
-        elif "face" in prompt_lower:
-            # Add facial details
-            center = (112, 112)
-            # Eyebrows
-            draw.arc([center[0]-25, center[1]-35, center[0]-5, center[1]-25], 0, 180, fill='black', width=2)
-            draw.arc([center[0]+5, center[1]-35, center[0]+25, center[1]-25], 0, 180, fill='black', width=2)
-            # Pupils
-            draw.ellipse([center[0]-17, center[1]-17, center[0]-13, center[1]-13], fill='black')
-            draw.ellipse([center[0]+13, center[1]-17, center[0]+17, center[1]-13], fill='black')
-        
-        return img
+        # Add a colored accent element
+        dwg.add(dwg.circle(
+            center=(width * 0.8, height * 0.2),
+            r=15,
+            fill=fill_color,
+            stroke='black',
+            stroke_width=1
+        ))
 
-    def _apply_style_transfer(self, img, prompt):
-        """Apply style transfer effects"""
-        prompt_lower = prompt.lower()
+    def _add_size_modifier(self, dwg, size_type, width, height):
+        """Add size modification indicators"""
+        if size_type == 'large':
+            # Add larger elements
+            dwg.add(dwg.rect(
+                insert=(10, 10),
+                size=(width-20, height-20),
+                fill='none',
+                stroke='gray',
+                stroke_width=3,
+                stroke_dasharray='5,5'
+            ))
+        elif size_type == 'small':
+            # Add smaller elements
+            dwg.add(dwg.rect(
+                insert=(width*0.3, height*0.3),
+                size=(width*0.4, height*0.4),
+                fill='none',
+                stroke='gray',
+                stroke_width=1,
+                stroke_dasharray='2,2'
+            ))
+
+    def _add_object_elements(self, dwg, obj_type, width, height):
+        """Add specific object elements"""
+        if obj_type == 'house':
+            self._add_house_elements(dwg, width, height)
+        elif obj_type == 'tree':
+            self._add_tree_elements(dwg, width, height)
+        elif obj_type == 'car':
+            self._add_car_elements(dwg, width, height)
+
+    def _add_detailed_elements(self, dwg, width, height, prompt):
+        """Add detailed elements for complex prompts"""
+        # Add multiple overlapping shapes for complexity
+        for i in range(8):
+            x = random.randint(20, width-40)
+            y = random.randint(20, height-40)
+            size = random.randint(10, 30)
+            
+            shape_type = random.choice(['circle', 'rect', 'polygon'])
+            
+            if shape_type == 'circle':
+                dwg.add(dwg.circle(
+                    center=(x, y),
+                    r=size,
+                    fill='none',
+                    stroke='black',
+                    stroke_width=1,
+                    opacity=0.7
+                ))
+            elif shape_type == 'rect':
+                dwg.add(dwg.rect(
+                    insert=(x-size, y-size),
+                    size=(size*2, size*2),
+                    fill='none',
+                    stroke='black',
+                    stroke_width=1,
+                    opacity=0.7
+                ))
+
+    def _add_simple_elements(self, dwg, width, height, prompt):
+        """Add simple elements for minimal prompts"""
+        # Add just a few basic shapes
+        center_x = width / 2
+        center_y = height / 2
         
-        if "cartoon" in prompt_lower:
-            # Cartoon style - enhance colors and add outlines
-            img = img.filter(ImageFilter.EDGE_ENHANCE_MORE)
-            # Increase saturation
-            from PIL import ImageEnhance
-            enhancer = ImageEnhance.Color(img)
-            img = enhancer.enhance(1.5)
-            
-        elif "realistic" in prompt_lower:
-            # Realistic style - add shading and texture
-            img = img.filter(ImageFilter.GaussianBlur(radius=0.5))
-            # Add subtle noise for texture
-            import numpy as np
-            pixels = np.array(img)
-            noise = np.random.normal(0, 3, pixels.shape)
-            pixels = np.clip(pixels + noise, 0, 255).astype(np.uint8)
-            img = Image.fromarray(pixels)
-            
-        elif "watercolor" in prompt_lower:
-            # Watercolor style - soft edges and bleeding
-            img = img.filter(ImageFilter.GaussianBlur(radius=1.0))
-            from PIL import ImageEnhance
-            enhancer = ImageEnhance.Color(img)
-            img = enhancer.enhance(0.8)
-            
-        return img
+        dwg.add(dwg.circle(
+            center=(center_x, center_y),
+            r=min(width, height) * 0.2,
+            fill='none',
+            stroke='black',
+            stroke_width=2
+        ))
 
-    def _add_details(self, img, prompt):
-        """Add contextual details to the sketch"""
-        draw = ImageDraw.Draw(img)
+    def _add_standard_elements(self, dwg, width, height, prompt):
+        """Add standard elements based on prompt"""
         prompt_lower = prompt.lower()
         
-        if "landscape" in prompt_lower:
-            # Add clouds
-            for i in range(3):
-                x = 40 + i * 60
-                y = 30 + i * 10
-                draw.ellipse([x, y, x+30, y+15], fill='lightgray', outline='gray')
-            # Add birds
-            for i in range(2):
-                x = 60 + i * 80
-                y = 50 + i * 5
-                draw.arc([x, y, x+10, y+5], 0, 180, fill='black', width=1)
-                
-        elif "portrait" in prompt_lower:
-            # Add hair details
-            center = (112, 112)
-            for i in range(15):
-                x = center[0] + random.randint(-50, 50)
-                y = center[1] - 60 + random.randint(0, 30)
-                draw.line([x, y, x + random.randint(-5, 5), y + random.randint(10, 30)], 
-                         fill='brown', width=1)
-                         
-        elif "building" in prompt_lower:
-            # Add architectural details
-            # Window frames
-            draw.rectangle([69, 139, 91, 161], outline='black', width=1)
-            draw.rectangle([133, 139, 155, 161], outline='black', width=1)
-            # Chimney
-            draw.rectangle([140, 70, 150, 90], outline='black', width=2)
-            
-        return img
+        if any(word in prompt_lower for word in ['house', 'building']):
+            self._add_house_elements(dwg, width, height)
+        elif any(word in prompt_lower for word in ['tree', 'forest']):
+            self._add_tree_elements(dwg, width, height)
+        elif any(word in prompt_lower for word in ['car', 'vehicle']):
+            self._add_car_elements(dwg, width, height)
+        else:
+            self._add_abstract_elements(dwg, width, height, prompt)
 
-    def _apply_coloring(self, img, prompt):
-        """Apply coloring to the sketch"""
-        # Convert to RGBA for transparency effects
-        img = img.convert('RGBA')
-        
-        # Create color overlay
-        color_overlay = Image.new('RGBA', img.size, (0, 0, 0, 0))
-        draw = ImageDraw.Draw(color_overlay)
-        
-        prompt_lower = prompt.lower()
+    def _add_abstract_elements(self, dwg, width, height, prompt):
+        """Add abstract elements based on prompt"""
+        prompt_hash = hash(prompt) % 100
         
-        if "sunset" in prompt_lower:
-            # Sunset colors
-            draw.rectangle([0, 0, 224, 112], fill=(255, 200, 100, 50))
-            draw.rectangle([0, 112, 224, 224], fill=(255, 150, 50, 30))
-            
-        elif "nature" in prompt_lower:
-            # Nature colors
-            draw.rectangle([0, 140, 224, 224], fill=(100, 200, 100, 40))  # Green ground
-            draw.rectangle([0, 0, 224, 140], fill=(150, 200, 255, 30))    # Blue sky
+        for i in range(5):
+            x = (i * 40 + prompt_hash) % (width - 40) + 20
+            y = (i * 35 + prompt_hash) % (height - 40) + 20
+            size = 15 + (i * 5) % 20
             
-        elif "warm" in prompt_lower:
-            # Warm color palette
-            draw.rectangle([0, 0, 224, 224], fill=(255, 200, 150, 25))
+            dwg.add(dwg.circle(
+                center=(x, y),
+                r=size,
+                fill='none',
+                stroke='black',
+                stroke_width=2,
+                opacity=0.8
+            ))
+
+    def _emphasize_element(self, dwg, word, weight, width, height):
+        """Emphasize an element based on attention weight"""
+        # Make elements larger and more prominent
+        scale_factor = weight
+        stroke_width = int(2 * scale_factor)
+        
+        if word in ['house', 'building']:
+            # Emphasized house
+            house_size = min(width, height) * 0.4 * scale_factor
+            house_x = (width - house_size) / 2
+            house_y = (height - house_size) / 2
             
-        elif "cool" in prompt_lower:
-            # Cool color palette
-            draw.rectangle([0, 0, 224, 224], fill=(150, 200, 255, 25))
+            dwg.add(dwg.rect(
+                insert=(house_x, house_y),
+                size=(house_size, house_size * 0.8),
+                fill='none',
+                stroke='red',
+                stroke_width=stroke_width
+            ))
+
+    def _deemphasize_element(self, dwg, word, weight, width, height):
+        """De-emphasize an element based on attention weight"""
+        # Make elements smaller and less prominent
+        scale_factor = weight
+        stroke_width = max(1, int(2 * scale_factor))
+        
+        if word in ['background', 'sky']:
+            # De-emphasized background elements
+            dwg.add(dwg.rect(
+                insert=(0, 0),
+                size=(width, height * 0.3),
+                fill='none',
+                stroke='lightgray',
+                stroke_width=stroke_width,
+                opacity=scale_factor
+            ))
+
+    def create_error_result(self, prompt: str, edit_type: str, error: str, width: int, height: int):
+        """Create error result with fallback SVG"""
+        fallback_svg = self.create_fallback_svg(prompt, width, height)
+        return {
+            "svg": fallback_svg,
+            "svg_base64": base64.b64encode(fallback_svg.encode('utf-8')).decode('utf-8'),
+            "edit_type": edit_type,
+            "prompt": prompt,
+            "error": error
+        }
+
+    def create_fallback_svg(self, prompt: str, width: int, height: int):
+        """Create simple fallback SVG"""
+        dwg = svgwrite.Drawing(size=(width, height))
+        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
         
-        # Blend with original
-        img = Image.alpha_composite(img, color_overlay)
-        img = img.convert('RGB')
+        # Simple centered text
+        dwg.add(dwg.text(
+            f"DiffSketchEdit\n{prompt[:20]}...",
+            insert=(width/2, height/2),
+            text_anchor="middle",
+            font_size="14",
+            fill="black"
+        ))
         
-        return img
+        return dwg.tostring()