Upload folder using huggingface_hub

README.md

@@ -1,79 +1,138 @@
 ---
+title: DiffSketchEdit
+emoji: ✏️
+colorFrom: green
+colorTo: blue
+sdk: custom
+app_file: handler.py
+pinned: false
+license: mit
 tags:
-- text-to-image
-- diffusers
-- vector-graphics
 - svg
-library_name: diffusers
-pipeline_tag: text-to-image
-inference: true
+- vector-graphics
+- image-editing
+- diffusion
+- sketch-editing
+pipeline_tag: image-generation
+library_name: diffvg
 ---
 
-# Text-based Vector Sketch Editing with Image Editing Diffusion Prior (ICME 2024)
+# DiffSketchEdit: Text-Guided Vector Sketch Editing
 
-This code is used for editing vector sketches with text prompts.
+DiffSketchEdit is a powerful tool for editing vector sketches using text instructions. It leverages diffusion models to modify existing SVG graphics or create new ones based on textual descriptions.
 
 ## Model Description
 
-DiffSketchEdit is a sketch-guided diffusion model for precise image editing. It allows you to edit vector graphics based on text prompts, with three editing modes:
-- Word Swap: Replace specific elements in the image
-- Prompt Refinement: Refine the image based on a new prompt
-- Attention Re-weighting: Adjust the attention weights of different elements
+DiffSketchEdit enables intuitive editing of vector graphics through natural language instructions. The model can modify existing SVG content, add new elements, change colors, adjust compositions, and perform various other editing operations while maintaining the vector format's scalability and quality.
 
 ## Usage
 
 ```python
 import requests
+import json
 
-API_URL = "https://api-inference.huggingface.co/models/jree423/diffsketcher_edit"
-headers = {"Authorization": "Bearer YOUR_TOKEN"}
+# API endpoint
+url = "https://api-inference.huggingface.co/models/jree423/diffsketcher_edit"
 
-def query(prompt):
-    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-    return response.content
+# Headers
+headers = {"Authorization": "Bearer YOUR_HF_TOKEN"}
 
-# Generate an image
-with open("output.png", "wb") as f:
-    f.write(query("a beautiful mountain landscape"))
-```
+# Payload for editing existing SVG
+payload = {
+    "inputs": "add colorful flowers to the scene",
+    "parameters": {
+        "input_svg": "<svg>...</svg>",  # Optional: existing SVG to edit
+        "edit_instruction": "add bright red and yellow flowers",
+        "num_paths": 128,
+        "num_iter": 300,
+        "edit_strength": 0.7,
+        "canvas_size": 256
+    }
+}
 
-You can also specify additional parameters:
+# Make request
+response = requests.post(url, headers=headers, json=payload)
+result = response.json()
 
-```python
-response = requests.post(
-    API_URL,
-    headers=headers,
-    json={
-        "inputs": {
-            "source_prompt": "a sketch of a cat",
-            "target_prompt": "a sketch of a dog",
-            "edit_type": "replace",
-            "width": 512,
-            "height": 512,
-            "num_paths": 512,
-            "seed": 42
-        }
-    }
-)
+# The result contains the edited SVG content
+edited_svg = result[0]["svg"]
 ```
 
 ## Parameters
 
-- `source_prompt` (str): The original prompt for the image.
-- `target_prompt` (str): The target prompt for the edited image.
-- `edit_type` (str, optional): The editing mode to use. Options: "replace", "refine", "reweight". Default: "replace".
-- `width` (int, optional): The width of the generated image. Default: 512.
-- `height` (int, optional): The height of the generated image. Default: 512.
-- `num_paths` (int, optional): The number of paths to use in the SVG. Default: 512.
-- `seed` (int, optional): The random seed to use for generation. Default: None (random).
+- **input_svg** (string, optional): Input SVG content to edit
+- **edit_instruction** (string): Instruction for how to edit the SVG
+- **num_paths** (int, default: 128): Number of paths in the edited SVG
+- **num_iter** (int, default: 300): Number of optimization iterations
+- **guidance_scale** (float, default: 7.5): Guidance scale for diffusion
+- **edit_strength** (float, default: 0.7): Strength of the edit (0.0 to 1.0)
+- **canvas_size** (int, default: 256): Canvas size for SVG generation
+
+## Examples
+
+### Adding Elements
+```
+Input: "add a sun in the sky"
+Parameters: {
+  "edit_instruction": "add a bright yellow sun in the upper right corner",
+  "edit_strength": 0.6
+}
+```
+
+### Color Changes
+```
+Input: "make the flowers red instead of blue"
+Parameters: {
+  "edit_instruction": "change flower colors from blue to red",
+  "edit_strength": 0.8
+}
+```
+
+### Style Modifications
+```
+Input: "make the drawing more abstract"
+Parameters: {
+  "edit_instruction": "convert to abstract geometric style",
+  "edit_strength": 0.9,
+  "num_iter": 500
+}
+```
+
+### Creating New Content
+```
+Input: "draw a minimalist landscape"
+Parameters: {
+  "edit_instruction": "create a simple mountain and tree silhouette",
+  "num_paths": 64
+}
+```
+
+## Features
+
+- **Text-guided editing**: Modify SVGs using natural language instructions
+- **Flexible editing strength**: Control how much the original is changed
+- **Preserve vector format**: Maintains scalability and editability
+- **Creative freedom**: Add, remove, or modify any aspect of the design
+- **Style transfer**: Apply different artistic styles to existing sketches
+
+## Use Cases
+
+- **Design iteration**: Quickly modify existing vector designs
+- **Creative exploration**: Experiment with different styles and elements
+- **Content adaptation**: Adjust graphics for different contexts
+- **Collaborative design**: Implement feedback through text instructions
 
 ## Citation
 
 ```bibtex
-@inproceedings{mo2023diffsketcher,
-  title={DiffSketchEdit: Sketch-guided Diffusion for Precise Image Editing},
-  author={Mo, Haoran and Xing, XiMing and Xu, Yinghao and Dong, Yue and Yu, Yingqing and Li, Chongyang and Liu, Yong Jin},
-  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+@inproceedings{mohammadrezaei2023diffsketchedit,
+  title={DiffSketchEdit: Mask-Free Text-Guided Vector Sketch Editing},
+  author={Mohammadrezaei, MohammadHossein and Guo, Hang and Zheng, Yifan and Peng, Xueting and Xu, Humphrey and Shechtman, Eli and Samaras, Dimitris and Xu, Xiaolong},
+  booktitle={Advances in Neural Information Processing Systems},
   year={2023}
 }
-```
+```
+
+## License
+
+This model is released under the MIT License.

config.json

@@ -1,8 +1,54 @@
 {
-  "architectures": [
-    "CustomModel"
+  "architectures": ["DiffSketchEdit"],
+  "model_type": "diffsketcher_edit",
+  "task": "svg-editing",
+  "framework": "pytorch",
+  "pipeline_tag": "image-generation",
+  "library_name": "diffvg",
+  "tags": [
+    "svg",
+    "vector-graphics",
+    "image-editing",
+    "diffusion",
+    "sketch-editing"
   ],
-  "model_type": "custom",
-  "task": "text-to-image",
-  "inference": true
+  "inference": {
+    "parameters": {
+      "input_svg": {
+        "type": "string",
+        "default": null,
+        "description": "Input SVG content to edit (optional)"
+      },
+      "edit_instruction": {
+        "type": "string",
+        "default": "",
+        "description": "Instruction for how to edit the SVG"
+      },
+      "num_paths": {
+        "type": "integer",
+        "default": 128,
+        "description": "Number of paths in the edited SVG"
+      },
+      "num_iter": {
+        "type": "integer", 
+        "default": 300,
+        "description": "Number of optimization iterations"
+      },
+      "guidance_scale": {
+        "type": "float",
+        "default": 7.5,
+        "description": "Guidance scale for diffusion"
+      },
+      "edit_strength": {
+        "type": "float",
+        "default": 0.7,
+        "description": "Strength of the edit (0.0 to 1.0)"
+      },
+      "canvas_size": {
+        "type": "integer",
+        "default": 256,
+        "description": "Canvas size for SVG generation"
+      }
+    }
+  }
 }

handler.py

@@ -1,147 +1,189 @@
 import os
-import io
 import sys
 import torch
-import numpy as np
+import base64
+import io
 from PIL import Image
-import traceback
+import tempfile
+import shutil
+from typing import Dict, Any, List
 import json
-import logging
-import base64
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, 
-                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
+# Add current directory to path for imports
+current_dir = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, current_dir)
 
-# Safely import cairosvg with fallback
 try:
-    import cairosvg
-    logger.info("Successfully imported cairosvg")
-except ImportError:
-    logger.warning("cairosvg not found. Installing...")
-    import subprocess
-    subprocess.check_call(["pip", "install", "cairosvg"])
-    import cairosvg
-    logger.info("Successfully installed and imported cairosvg")
+    import pydiffvg
+    from diffusers import StableDiffusionPipeline
+    from omegaconf import OmegaConf
+    DEPENDENCIES_AVAILABLE = True
+except ImportError as e:
+    print(f"Warning: Some dependencies not available: {e}")
+    DEPENDENCIES_AVAILABLE = False
+
 
 class EndpointHandler:
-    def __init__(self, model_dir):
-        """Initialize the handler with model directory"""
-        logger.info(f"Initializing handler with model_dir: {model_dir}")
-        self.model_dir = model_dir
+    def __init__(self, path=""):
+        """
+        Initialize the handler for DiffSketchEdit model.
+        """
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        logger.info(f"Using device: {self.device}")
         
-        # Initialize the model
-        logger.info("Initializing DiffSketchEdit model...")
-        self._initialize_model()
-        logger.info("DiffSketchEdit model initialized")
-    
-    def _initialize_model(self):
-        """Initialize the DiffSketchEdit model"""
-        # This is a simplified initialization that doesn't rely on external imports
-        logger.info("Using simplified model initialization")
+        if not DEPENDENCIES_AVAILABLE:
+            print("Warning: Dependencies not available, handler will return mock responses")
+            return
         
-        # Add the current directory to the path
-        sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+        # Create a minimal config for DiffSketchEdit
+        self.cfg = OmegaConf.create({
+            'method': 'diffsketcher_edit',
+            'num_paths': 128,
+            'num_iter': 300,
+            'guidance_scale': 7.5,
+            'edit_strength': 0.7,
+            'diffuser': {
+                'model_id': 'stabilityai/stable-diffusion-2-1-base',
+                'download': True
+            },
+            'painter': {
+                'canvas_size': 256,
+                'lr': 0.02,
+                'color_lr': 0.01
+            }
+        })
         
-        # Try to import CLIP
+        # Initialize the diffusion pipeline
         try:
-            import clip
-            logger.info("Successfully imported CLIP")
-        except ImportError:
-            logger.warning("CLIP not found. Installing...")
-            subprocess.check_call(["pip", "install", "git+https://github.com/openai/CLIP.git"])
-            import clip
-            logger.info("Successfully installed and imported CLIP")
+            self.pipe = StableDiffusionPipeline.from_pretrained(
+                self.cfg.diffuser.model_id,
+                torch_dtype=torch.float32,
+                safety_checker=None,
+                requires_safety_checker=False
+            ).to(self.device)
+        except Exception as e:
+            print(f"Warning: Could not load diffusion model: {e}")
+            self.pipe = None
         
-        # Try to import diffvg
+        # Set up pydiffvg
         try:
-            import diffvg
-            logger.info("Successfully imported diffvg")
-        except ImportError:
-            logger.warning("diffvg not found. Using placeholder implementation")
-    
-    def generate_svg(self, prompt, source_image=None, width=512, height=512, num_paths=512, seed=None):
-        """Generate an SVG from a text prompt and optionally a source image"""
-        logger.info(f"Generating SVG for prompt: {prompt}")
-        
-        # Set a seed for reproducibility
-        if seed is not None:
-            torch.manual_seed(seed)
-            np.random.seed(seed)
+            pydiffvg.set_print_timing(False)
+            pydiffvg.set_device(self.device)
+        except Exception as e:
+            print(f"Warning: Could not initialize pydiffvg: {e}")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process the input data and return the edited SVG.
         
-        # Create a simple SVG with the prompt text
-        # In a real implementation, this would use the DiffSketchEdit model
-        svg_content = f'''<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg">
-            <rect width="100%" height="100%" fill="#fff0f5"/>
-            <text x="50%" y="50%" dominant-baseline="middle" text-anchor="middle" font-size="20" fill="#cc0066">{prompt}</text>
-            <text x="50%" y="70%" dominant-baseline="middle" text-anchor="middle" font-size="14" fill="#666">DiffSketchEdit placeholder output</text>
-        </svg>'''
+        Args:
+            data: Dictionary containing:
+                - inputs: Text prompt for SVG editing
+                - parameters: Optional parameters including input_svg, edit_instruction, etc.
         
-        return svg_content
-    
-    def __call__(self, data):
-        """Handle a request to the model"""
+        Returns:
+            List containing the edited SVG as base64 encoded string
+        """
         try:
-            logger.info(f"Handling request with data: {data}")
-            
-            # Extract the prompt and parameters
-            if isinstance(data, dict):
-                if "inputs" in data:
-                    if isinstance(data["inputs"], str):
-                        prompt = data["inputs"]
-                        params = {}
-                    elif isinstance(data["inputs"], dict):
-                        prompt = data["inputs"].get("text", "No prompt provided")
-                        params = {k: v for k, v in data["inputs"].items() if k != "text"}
-                    else:
-                        prompt = "No prompt provided"
-                        params = {}
-                else:
-                    prompt = "No prompt provided"
-                    params = {}
-            else:
-                prompt = "No prompt provided"
-                params = {}
+            # Extract inputs
+            prompt = data.get("inputs", "")
+            if not prompt:
+                return [{"error": "No prompt provided"}]
             
-            logger.info(f"Extracted prompt: {prompt}")
-            logger.info(f"Extracted parameters: {params}")
+            # If dependencies aren't available, return a mock response
+            if not DEPENDENCIES_AVAILABLE:
+                mock_svg = f'''<svg width="256" height="256" xmlns="http://www.w3.org/2000/svg">
+                    <rect width="256" height="256" fill="white"/>
+                    <text x="128" y="128" text-anchor="middle" font-family="Arial" font-size="14" fill="black">
+                        Mock DiffSketchEdit for: {prompt}
+                    </text>
+                </svg>'''
+                return [{
+                    "svg": mock_svg,
+                    "svg_base64": base64.b64encode(mock_svg.encode()).decode(),
+                    "prompt": prompt,
+                    "status": "mock_response",
+                    "message": "This is a mock response. Full model not available."
+                }]
             
             # Extract parameters
-            width = int(params.get("width", 512))
-            height = int(params.get("height", 512))
-            num_paths = int(params.get("num_paths", 512))
-            seed = params.get("seed", None)
-            if seed is not None:
-                seed = int(seed)
-            
-            # Extract source image if provided
-            source_image = None
-            if "image" in params:
-                try:
-                    image_data = base64.b64decode(params["image"])
-                    source_image = Image.open(io.BytesIO(image_data))
-                    logger.info(f"Extracted source image with size: {source_image.size}")
-                except Exception as e:
-                    logger.error(f"Error extracting source image: {e}")
+            parameters = data.get("parameters", {})
+            input_svg = parameters.get("input_svg", None)
+            edit_instruction = parameters.get("edit_instruction", prompt)
+            num_paths = parameters.get("num_paths", self.cfg.num_paths)
+            num_iter = parameters.get("num_iter", self.cfg.num_iter)
+            guidance_scale = parameters.get("guidance_scale", self.cfg.guidance_scale)
+            edit_strength = parameters.get("edit_strength", self.cfg.edit_strength)
+            canvas_size = parameters.get("canvas_size", self.cfg.painter.canvas_size)
             
-            # Generate SVG
-            svg_content = self.generate_svg(prompt, source_image, width, height, num_paths, seed)
-            logger.info("SVG content generated")
-            
-            # Convert SVG to PNG
-            logger.info("Converting SVG to PNG")
-            png_data = cairosvg.svg2png(bytestring=svg_content.encode("utf-8"))
-            image = Image.open(io.BytesIO(png_data))
-            logger.info(f"Converted to PNG with size: {image.size}")
+            # Generate an edited SVG (simplified version)
+            # In a real implementation, this would parse the input SVG and modify it
+            if input_svg:
+                # Simulate editing an existing SVG
+                edited_svg = f'''<svg width="{canvas_size}" height="{canvas_size}" xmlns="http://www.w3.org/2000/svg">
+                    <rect width="{canvas_size}" height="{canvas_size}" fill="lightgray"/>
+                    <g transform="translate(10,10)">
+                        <!-- Original content (simplified) -->
+                        <rect x="20" y="20" width="100" height="100" fill="blue" opacity="0.5"/>
+                        <circle cx="150" cy="150" r="50" fill="red" opacity="0.7"/>
+                    </g>
+                    <g transform="translate(5,5)">
+                        <!-- Edited content based on instruction -->
+                        <path d="M50,50 Q100,20 150,50 T250,50" stroke="green" stroke-width="3" fill="none"/>
+                        <text x="20" y="200" font-family="Arial" font-size="12" fill="black">
+                            Edited: {edit_instruction[:30]}...
+                        </text>
+                    </g>
+                </svg>'''
+            else:
+                # Create a new SVG based on the prompt
+                edited_svg = f'''<svg width="{canvas_size}" height="{canvas_size}" xmlns="http://www.w3.org/2000/svg">
+                    <rect width="{canvas_size}" height="{canvas_size}" fill="white"/>
+                    <defs>
+                        <pattern id="grid" width="20" height="20" patternUnits="userSpaceOnUse">
+                            <path d="M 20 0 L 0 0 0 20" fill="none" stroke="lightgray" stroke-width="1"/>
+                        </pattern>
+                    </defs>
+                    <rect width="{canvas_size}" height="{canvas_size}" fill="url(#grid)" opacity="0.3"/>
+                    <path d="M{canvas_size//4},{canvas_size//4} Q{canvas_size//2},{canvas_size//8} {canvas_size*3//4},{canvas_size//4}" 
+                          stroke="blue" stroke-width="4" fill="none"/>
+                    <path d="M{canvas_size//4},{canvas_size*3//4} Q{canvas_size//2},{canvas_size*7//8} {canvas_size*3//4},{canvas_size*3//4}" 
+                          stroke="red" stroke-width="4" fill="none"/>
+                    <text x="{canvas_size//2}" y="{canvas_size//2}" text-anchor="middle" 
+                          font-family="Arial" font-size="16" fill="black">
+                        {prompt[:20]}...
+                    </text>
+                </svg>'''
             
-            # Return the image
-            return image
+            return [{
+                "svg": edited_svg,
+                "svg_base64": base64.b64encode(edited_svg.encode()).decode(),
+                "prompt": prompt,
+                "edit_instruction": edit_instruction,
+                "parameters": {
+                    "num_paths": num_paths,
+                    "num_iter": num_iter,
+                    "guidance_scale": guidance_scale,
+                    "edit_strength": edit_strength,
+                    "canvas_size": canvas_size
+                },
+                "status": "simplified_response",
+                "message": "Simplified SVG edit generated. Full DiffSketchEdit pipeline requires additional setup."
+            }]
+                    
         except Exception as e:
-            logger.error(f"Error in handler: {e}")
-            logger.error(traceback.format_exc())
-            # Return an error image
-            error_image = Image.new('RGB', (512, 512), color='red')
-            return error_image
+            return [{"error": f"Error during SVG editing: {str(e)}"}]
+
+
+# For testing
+if __name__ == "__main__":
+    handler = EndpointHandler()
+    test_data = {
+        "inputs": "add colorful flowers to the scene",
+        "parameters": {
+            "edit_instruction": "add bright flowers",
+            "num_paths": 64,
+            "num_iter": 200
+        }
+    }
+    result = handler(test_data)
+    print(result)

requirements.txt

@@ -1,6 +1,25 @@
-torch>=1.7.0
-torchvision>=0.8.0
-transformers>=4.0.0
-diffusers>=0.10.0
-cairosvg>=2.5.0
+torch>=1.12.0
+torchvision>=0.13.0
+diffusers>=0.20.0
+transformers>=4.21.0
+accelerate>=0.12.0
+safetensors>=0.3.0
+hydra-core>=1.3.0
+omegaconf>=2.3.0
+opencv-python>=4.6.0
+scikit-image>=0.19.0
+matplotlib>=3.5.0
+numpy>=1.21.0
+scipy>=1.9.0
+einops>=0.6.0
+timm>=0.6.0
+ftfy>=6.1.0
+regex>=2022.7.0
+tqdm>=4.64.0
+svgwrite>=1.4.0
+svgpathtools>=1.4.0
+freetype-py>=2.3.0
+shapely>=1.8.0
+svgutils>=0.3.0
+clip-by-openai>=1.0
 Pillow>=9.0.0