Upload handler.py with huggingface_hub

handler.py

@@ -0,0 +1,109 @@
+import os
+import io
+import sys
+import torch
+import numpy as np
+from PIL import Image
+import traceback
+import json
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, 
+                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Add the model directory to the path
+sys.path.append('/code/diffsketcher_edit')
+
+# 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")
+
+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
+        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 diffsketcher_edit model...")
+        self.model = self._initialize_model()
+        logger.info("diffsketcher_edit model initialized")
+    
+    def _initialize_model(self):
+        # Initialize the diffsketcher_edit model
+        # This is a placeholder for the actual model initialization
+        # In a real implementation, you would load the model weights and initialize the model
+        return None
+    
+    def generate_svg(self, prompt, width=512, height=512, num_paths=512, seed=None):
+        # Generate an SVG from a text prompt
+        logger.info(f"Generating SVG for prompt: {prompt}")
+        
+        # This is a placeholder for the actual SVG generation
+        # In a real implementation, you would use the model to generate an SVG
+        svg_content = f'<svg width="{width}" height="{height}" xmlns="http://www.w3.org/2000/svg"><text x="50%" y="50%" dominant-baseline="middle" text-anchor="middle" font-size="20">{prompt}</text></svg>'
+        
+        return svg_content
+    
+    def __call__(self, data):
+        # Handle a request to the model
+        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 = {}
+            
+            logger.info(f"Extracted prompt: {prompt}")
+            logger.info(f"Extracted parameters: {params}")
+            
+            # Extract parameters
+            width = params.get("width", 512)
+            height = params.get("height", 512)
+            num_paths = params.get("num_paths", 512)
+            seed = params.get("seed", None)
+            
+            # Generate SVG
+            svg_content = self.generate_svg(prompt, 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}")
+            
+            # Return the PIL Image directly
+            return image
+        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