Upload handler.py with huggingface_hub

handler.py

@@ -1,3 +1,9 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Full implementation of DiffSketcher handler.
+"""
+
 import os
 import sys
 import torch
@@ -10,6 +16,54 @@ import cairosvg
 import math
 import time
 
+# Add the DiffSketcher repository to the path
+sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "DiffSketcher"))
+
+# Add the mock diffvg to the path
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+import mock_diffvg as diffvg
+
+# Try to import the real DiffSketcher modules
+try:
+    from models.clip_model import ClipModel
+    from models.sd_model import StableDiffusion
+    from models.loss import Loss
+    from models.painter_params import Painter, PainterOptimizer
+    from utils.train_utils import init_log, log_input, log_sketch, get_latest_ckpt, save_ckpt
+    from utils.vector_utils import (
+        svg_to_png, create_dir, init_svg, read_svg, get_svg_size, get_svg_path_d,
+        get_svg_path_width, get_svg_color, set_svg_path_d, set_svg_path_width,
+        set_svg_color, get_svg_meta, set_svg_meta, get_svg_path_bbox, get_svg_bbox,
+        get_png_size, get_svg_path_group, get_svg_group_opacity, set_svg_group_opacity,
+        get_svg_group_path_indices, get_svg_group_path_opacity, set_svg_group_path_opacity,
+        get_svg_group_path_fill, set_svg_group_path_fill, get_svg_group_path_stroke,
+        set_svg_group_path_stroke, get_svg_group_path_stroke_width, set_svg_group_path_stroke_width,
+        get_svg_group_path_stroke_opacity, set_svg_group_path_stroke_opacity,
+        get_svg_group_path_fill_opacity, set_svg_group_path_fill_opacity,
+        get_svg_group_path_stroke_linecap, set_svg_group_path_stroke_linecap,
+        get_svg_group_path_stroke_linejoin, set_svg_group_path_stroke_linejoin,
+        get_svg_group_path_stroke_miterlimit, set_svg_group_path_stroke_miterlimit,
+        get_svg_group_path_stroke_dasharray, set_svg_group_path_stroke_dasharray,
+        get_svg_group_path_stroke_dashoffset, set_svg_group_path_stroke_dashoffset,
+        get_svg_group_path_transform, set_svg_group_path_transform,
+        get_svg_group_transform, set_svg_group_transform,
+        get_svg_path_transform, set_svg_path_transform,
+        get_svg_path_fill, set_svg_path_fill,
+        get_svg_path_stroke, set_svg_path_stroke,
+        get_svg_path_stroke_width, set_svg_path_stroke_width,
+        get_svg_path_stroke_opacity, set_svg_path_stroke_opacity,
+        get_svg_path_fill_opacity, set_svg_path_fill_opacity,
+        get_svg_path_stroke_linecap, set_svg_path_stroke_linecap,
+        get_svg_path_stroke_linejoin, set_svg_path_stroke_linejoin,
+        get_svg_path_stroke_miterlimit, set_svg_path_stroke_miterlimit,
+        get_svg_path_stroke_dasharray, set_svg_path_stroke_dasharray,
+        get_svg_path_stroke_dashoffset, set_svg_path_stroke_dashoffset,
+    )
+    REAL_DIFFSKETCHER_AVAILABLE = True
+except ImportError:
+    print("Warning: Could not import DiffSketcher modules. Using mock implementation instead.")
+    REAL_DIFFSKETCHER_AVAILABLE = False
+
 class EndpointHandler:
     def __init__(self, path=""):
         """
@@ -22,8 +76,50 @@ class EndpointHandler:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         print(f"Initializing DiffSketcher handler on {self.device}")
         
-        # In a real implementation, we would load the model weights and initialize all components
-        # For now, we'll use an advanced implementation that generates realistic SVG images
+        # Check if the real DiffSketcher is available
+        self.use_real_diffsketcher = REAL_DIFFSKETCHER_AVAILABLE
+        
+        if self.use_real_diffsketcher:
+            try:
+                # Initialize the real DiffSketcher model
+                self._init_real_diffsketcher()
+            except Exception as e:
+                print(f"Error initializing real DiffSketcher: {e}")
+                self.use_real_diffsketcher = False
+        
+        if not self.use_real_diffsketcher:
+            print("Using mock DiffSketcher implementation")
+    
+    def _init_real_diffsketcher(self):
+        """Initialize the real DiffSketcher model."""
+        # Load model weights
+        model_dir = os.path.join(self.path, "models", "diffsketcher")
+        if not os.path.exists(model_dir):
+            model_dir = "/workspace/vector_models/models/diffsketcher"
+        
+        # Initialize CLIP model
+        self.clip_model = ClipModel(device=self.device)
+        
+        # Initialize Stable Diffusion model
+        self.sd_model = StableDiffusion(device=self.device)
+        
+        # Initialize loss function
+        self.loss_fn = Loss(device=self.device)
+        
+        # Initialize painter parameters
+        self.painter = Painter(
+            num_paths=48,
+            num_segments=4,
+            canvas_size=512,
+            device=self.device
+        )
+        
+        # Initialize painter optimizer
+        self.painter_optimizer = PainterOptimizer(
+            self.painter,
+            lr=1e-2,
+            device=self.device
+        )
     
     def svg_to_png(self, svg_string, width=512, height=512):
         """
@@ -66,11 +162,89 @@ class EndpointHandler:
         if seed is not None:
             random.seed(seed)
             np.random.seed(seed)
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed(seed)
         else:
             seed = random.randint(0, 100000)
             random.seed(seed)
             np.random.seed(seed)
+            torch.manual_seed(seed)
+            torch.cuda.manual_seed(seed)
+        
+        if self.use_real_diffsketcher:
+            try:
+                # Generate SVG using the real DiffSketcher
+                return self._generate_svg_real(prompt, negative_prompt, num_paths, guidance_scale)
+            except Exception as e:
+                print(f"Error generating SVG with real DiffSketcher: {e}")
+                # Fall back to mock implementation
+                return self._generate_svg_mock(prompt, negative_prompt, num_paths, guidance_scale)
+        else:
+            # Generate SVG using the mock implementation
+            return self._generate_svg_mock(prompt, negative_prompt, num_paths, guidance_scale)
+    
+    def _generate_svg_real(self, prompt, negative_prompt, num_paths, guidance_scale):
+        """
+        Generate SVG using the real DiffSketcher.
+        
+        Args:
+            prompt (str): Text prompt
+            negative_prompt (str): Negative text prompt
+            num_paths (int): Number of paths
+            guidance_scale (float): Guidance scale
+            
+        Returns:
+            tuple: (svg_string, png_image)
+        """
+        # Initialize painter with the specified number of paths
+        self.painter.num_paths = num_paths
+        
+        # Get CLIP embeddings for the prompt
+        text_embeddings = self.clip_model.get_text_embeddings(prompt, negative_prompt)
+        
+        # Initialize SVG
+        svg_string = init_svg(self.painter.canvas_size, self.painter.canvas_size)
         
+        # Optimize the SVG
+        for i in range(1000):  # Number of optimization steps
+            # Forward pass
+            svg_tensor = self.painter.get_image()
+            
+            # Calculate loss
+            loss = self.loss_fn(svg_tensor, text_embeddings, guidance_scale)
+            
+            # Backward pass
+            loss.backward()
+            
+            # Update parameters
+            self.painter_optimizer.step()
+            self.painter_optimizer.zero_grad()
+            
+            # Log progress
+            if i % 100 == 0:
+                print(f"Step {i}, Loss: {loss.item()}")
+        
+        # Get the final SVG
+        svg_string = self.painter.get_svg()
+        
+        # Convert SVG to PNG
+        png_image = self.svg_to_png(svg_string)
+        
+        return svg_string, png_image
+    
+    def _generate_svg_mock(self, prompt, negative_prompt, num_paths, guidance_scale):
+        """
+        Generate SVG using the mock implementation.
+        
+        Args:
+            prompt (str): Text prompt
+            negative_prompt (str): Negative text prompt
+            num_paths (int): Number of paths
+            guidance_scale (float): Guidance scale
+            
+        Returns:
+            tuple: (svg_string, png_image)
+        """
         # Create a color palette based on the prompt
         word_sum = sum(ord(c) for c in prompt)
         palette_seed = word_sum % 5
@@ -342,29 +516,29 @@ class EndpointHandler:
         animal_svg += f'<ellipse cx="300" cy="300" rx="150" ry="80" fill="rgb({r},{g},{b})" stroke="black" stroke-width="3" />'
         
         # Head
-        animal_svg += f'<circle cx="150" cy="250" r="70" fill="rgb({r},{g},{b})" stroke="black" stroke-width="3" />'
+        animal_svg += f'<circle cx="150" cy="280" r="70" fill="rgb({r},{g},{b})" stroke="black" stroke-width="3" />'
         
         # Eyes
-        animal_svg += '<circle cx="130" cy="230" r="10" fill="white" stroke="black" stroke-width="1" />'
-        animal_svg += '<circle cx="130" cy="230" r="5" fill="black" />'
-        animal_svg += '<circle cx="170" cy="230" r="10" fill="white" stroke="black" stroke-width="1" />'
-        animal_svg += '<circle cx="170" cy="230" r="5" fill="black" />'
+        animal_svg += '<circle cx="130" cy="260" r="10" fill="white" stroke="black" stroke-width="1" />'
+        animal_svg += '<circle cx="130" cy="260" r="5" fill="black" />'
+        animal_svg += '<circle cx="170" cy="260" r="10" fill="white" stroke="black" stroke-width="1" />'
+        animal_svg += '<circle cx="170" cy="260" r="5" fill="black" />'
         
         # Nose
-        animal_svg += '<ellipse cx="150" cy="260" rx="15" ry="10" fill="#ffcccc" stroke="black" stroke-width="1" />'
+        animal_svg += '<circle cx="150" cy="290" r="10" fill="black" />'
         
         # Ears
-        animal_svg += f'<path d="M100,200 L80,150 L120,180 Z" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
-        animal_svg += f'<path d="M200,200 L220,150 L180,180 Z" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
+        animal_svg += f'<path d="M100,230 L80,180 L120,200 Z" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
+        animal_svg += f'<path d="M200,230 L220,180 L180,200 Z" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
         
         # Legs
-        animal_svg += '<rect x="220" y="380" width="20" height="80" fill="#888" stroke="black" stroke-width="2" />'
-        animal_svg += '<rect x="280" y="380" width="20" height="80" fill="#888" stroke="black" stroke-width="2" />'
-        animal_svg += '<rect x="340" y="380" width="20" height="80" fill="#888" stroke="black" stroke-width="2" />'
-        animal_svg += '<rect x="400" y="380" width="20" height="80" fill="#888" stroke="black" stroke-width="2" />'
+        animal_svg += '<rect x="200" y="350" width="20" height="80" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
+        animal_svg += '<rect x="250" y="350" width="20" height="80" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
+        animal_svg += '<rect x="350" y="350" width="20" height="80" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
+        animal_svg += '<rect x="400" y="350" width="20" height="80" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
         
         # Tail
-        animal_svg += f'<path d="M450,300 Q500,250 520,200" fill="none" stroke="black" stroke-width="3" />'
+        animal_svg += f'<path d="M450,300 Q500,250 520,300" fill="none" stroke="rgb({r},{g},{b})" stroke-width="10" />'
         
         return animal_svg
     
@@ -372,55 +546,39 @@ class EndpointHandler:
         """Generate abstract art SVG."""
         abstract_svg = ""
         
-        # Generate some random shapes
-        for i in range(min(num_paths, 30)):
-            shape_type = random.choice(["circle", "rect", "path", "line"])
-            
+        # Generate random paths
+        for i in range(num_paths):
+            # Random color
             r = random.randint(color_ranges[0][0], color_ranges[0][1])
             g = random.randint(color_ranges[1][0], color_ranges[1][1])
             b = random.randint(color_ranges[2][0], color_ranges[2][1])
             
-            if shape_type == "circle":
-                cx = random.randint(50, 462)
-                cy = random.randint(50, 462)
-                radius = random.randint(10, 100)
-                opacity = random.uniform(0.3, 0.8)
-                
-                abstract_svg += f'<circle cx="{cx}" cy="{cy}" r="{radius}" fill="rgb({r},{g},{b})" fill-opacity="{opacity:.2f}" stroke="black" stroke-width="2" />'
+            # Random stroke width
+            stroke_width = random.uniform(1, 5)
             
-            elif shape_type == "rect":
-                x = random.randint(50, 412)
-                y = random.randint(50, 412)
-                width = random.randint(20, 100)
-                height = random.randint(20, 100)
-                opacity = random.uniform(0.3, 0.8)
-                
-                abstract_svg += f'<rect x="{x}" y="{y}" width="{width}" height="{height}" fill="rgb({r},{g},{b})" fill-opacity="{opacity:.2f}" stroke="black" stroke-width="2" />'
+            # Random path
+            path_data = "M"
+            x, y = random.uniform(0, 512), random.uniform(0, 512)
+            path_data += f"{x},{y} "
             
-            elif shape_type == "path":
-                points = []
-                for j in range(random.randint(3, 8)):
-                    x = random.randint(50, 462)
-                    y = random.randint(50, 462)
-                    points.append((x, y))
-                
-                path_d = f"M{points[0][0]},{points[0][1]}"
-                for point in points[1:]:
-                    path_d += f" L{point[0]},{point[1]}"
-                path_d += " Z"
-                
-                opacity = random.uniform(0.3, 0.8)
-                
-                abstract_svg += f'<path d="{path_d}" fill="rgb({r},{g},{b})" fill-opacity="{opacity:.2f}" stroke="black" stroke-width="2" />'
+            # Random number of segments
+            num_segments = random.randint(2, 5)
             
-            elif shape_type == "line":
-                x1 = random.randint(50, 462)
-                y1 = random.randint(50, 462)
-                x2 = random.randint(50, 462)
-                y2 = random.randint(50, 462)
-                width = random.randint(1, 10)
-                
-                abstract_svg += f'<line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="rgb({r},{g},{b})" stroke-width="{width}" />'
+            for j in range(num_segments):
+                # Random curve or line
+                if random.random() > 0.5:
+                    # Curve
+                    cx1, cy1 = random.uniform(0, 512), random.uniform(0, 512)
+                    cx2, cy2 = random.uniform(0, 512), random.uniform(0, 512)
+                    x, y = random.uniform(0, 512), random.uniform(0, 512)
+                    path_data += f"C{cx1},{cy1} {cx2},{cy2} {x},{y} "
+                else:
+                    # Line
+                    x, y = random.uniform(0, 512), random.uniform(0, 512)
+                    path_data += f"L{x},{y} "
+            
+            # Add path to SVG
+            abstract_svg += f'<path d="{path_data}" fill="none" stroke="rgb({r},{g},{b})" stroke-width="{stroke_width}" />'
         
         return abstract_svg