Upload handler.py with huggingface_hub

handler.py

@@ -7,15 +7,13 @@ from PIL import Image
 import traceback
 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 the model directory to the path
-sys.path.append('/code/diffsketcher_edit')
-
 # Safely import cairosvg with fallback
 try:
     import cairosvg
@@ -29,35 +27,63 @@ except ImportError:
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        # Initialize the handler with model directory
+        """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")
+        logger.info("Initializing DiffSketchEdit model...")
+        self._initialize_model()
+        logger.info("DiffSketchEdit 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
+        """Initialize the DiffSketchEdit model"""
+        # This is a simplified initialization that doesn't rely on external imports
+        logger.info("Using simplified model initialization")
+        
+        # Add the current directory to the path
+        sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+        
+        # Try to import CLIP
+        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")
+        
+        # Try to import diffvg
+        try:
+            import diffvg
+            logger.info("Successfully imported diffvg")
+        except ImportError:
+            logger.warning("diffvg not found. Using placeholder implementation")
     
-    def generate_svg(self, prompt, width=512, height=512, num_paths=512, seed=None):
-        # Generate an SVG from a text prompt
+    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}")
         
-        # 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>'
+        # Set a seed for reproducibility
+        if seed is not None:
+            torch.manual_seed(seed)
+            np.random.seed(seed)
+        
+        # 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>'''
         
         return svg_content
     
     def __call__(self, data):
-        # Handle a request to the model
+        """Handle a request to the model"""
         try:
             logger.info(f"Handling request with data: {data}")
             
@@ -84,13 +110,25 @@ class EndpointHandler:
             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)
+            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}")
             
             # Generate SVG
-            svg_content = self.generate_svg(prompt, width, height, num_paths, seed)
+            svg_content = self.generate_svg(prompt, source_image, width, height, num_paths, seed)
             logger.info("SVG content generated")
             
             # Convert SVG to PNG
@@ -99,11 +137,11 @@ class EndpointHandler:
             image = Image.open(io.BytesIO(png_data))
             logger.info(f"Converted to PNG with size: {image.size}")
             
-            # Return the PIL Image directly
+            # Return the image
             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
+            return error_image