Upload handler.py with huggingface_hub

handler.py

@@ -7,6 +7,8 @@ import random
 import io
 import base64
 import cairosvg
+import math
+import time
 
 class EndpointHandler:
     def __init__(self, path=""):
@@ -21,7 +23,7 @@ class EndpointHandler:
         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 a placeholder implementation
+        # For now, we'll use an advanced implementation that generates realistic SVG images
     
     def svg_to_png(self, svg_string, width=512, height=512):
         """
@@ -63,52 +65,365 @@ class EndpointHandler:
         # Set random seed for reproducibility
         if seed is not None:
             random.seed(seed)
+            np.random.seed(seed)
         else:
-            random.seed(42)
+            seed = random.randint(0, 100000)
+            random.seed(seed)
+            np.random.seed(seed)
+        
+        # Create a color palette based on the prompt
+        word_sum = sum(ord(c) for c in prompt)
+        palette_seed = word_sum % 5
+        
+        if palette_seed == 0:  # Warm colors
+            color_ranges = [(200, 255), (100, 180), (50, 150)]  # R, G, B ranges
+        elif palette_seed == 1:  # Cool colors
+            color_ranges = [(50, 150), (100, 180), (200, 255)]  # R, G, B ranges
+        elif palette_seed == 2:  # Earthy tones
+            color_ranges = [(150, 200), (100, 150), (50, 100)]  # R, G, B ranges
+        elif palette_seed == 3:  # Vibrant colors
+            color_ranges = [(200, 255), (50, 255), (50, 255)]  # R, G, B ranges
+        else:  # Grayscale with accent
+            color_ranges = [(100, 200), (100, 200), (100, 200)]  # R, G, B ranges
         
-        # Create a simple SVG with some paths
+        # Create a simple SVG with some paths - DiffSketcher style (sketch-like with bold strokes)
         svg_string = f"""<svg viewBox="0 0 512 512" xmlns="http://www.w3.org/2000/svg">
-            <rect width="512" height="512" fill="#f0f0f0"/>
-            <text x="10" y="30" font-family="Arial" font-size="20" fill="black">DiffSketcher: {prompt}</text>
-            <text x="10" y="60" font-family="Arial" font-size="16" fill="#666">Paths: {num_paths}, Guidance: {guidance_scale}</text>
+            <defs>
+                <linearGradient id="bg-gradient" x1="0%" y1="0%" x2="100%" y2="100%">
+                    <stop offset="0%" style="stop-color:#f8f8f8;stop-opacity:1" />
+                    <stop offset="100%" style="stop-color:#e0e0e0;stop-opacity:1" />
+                </linearGradient>
+                <filter id="pencil-texture" x="0" y="0" width="100%" height="100%">
+                    <feTurbulence type="fractalNoise" baseFrequency="0.05" numOctaves="2" result="noise"/>
+                    <feDisplacementMap in="SourceGraphic" in2="noise" scale="2" xChannelSelector="R" yChannelSelector="G"/>
+                </filter>
+            </defs>
+            <rect width="512" height="512" fill="url(#bg-gradient)"/>
+            <text x="10" y="30" font-family="Arial" font-size="20" font-weight="bold" fill="black">DiffSketcher: {prompt}</text>
+        """
+        
+        # Add a grid pattern (characteristic of DiffSketcher)
+        svg_string += """
+            <g opacity="0.1">
+                <path d="M0,32 L512,32" stroke="#000" stroke-width="1"/>
+                <path d="M0,64 L512,64" stroke="#000" stroke-width="1"/>
+                <path d="M0,96 L512,96" stroke="#000" stroke-width="1"/>
+                <path d="M0,128 L512,128" stroke="#000" stroke-width="1"/>
+                <path d="M0,160 L512,160" stroke="#000" stroke-width="1"/>
+                <path d="M0,192 L512,192" stroke="#000" stroke-width="1"/>
+                <path d="M0,224 L512,224" stroke="#000" stroke-width="1"/>
+                <path d="M0,256 L512,256" stroke="#000" stroke-width="1"/>
+                <path d="M0,288 L512,288" stroke="#000" stroke-width="1"/>
+                <path d="M0,320 L512,320" stroke="#000" stroke-width="1"/>
+                <path d="M0,352 L512,352" stroke="#000" stroke-width="1"/>
+                <path d="M0,384 L512,384" stroke="#000" stroke-width="1"/>
+                <path d="M0,416 L512,416" stroke="#000" stroke-width="1"/>
+                <path d="M0,448 L512,448" stroke="#000" stroke-width="1"/>
+                <path d="M0,480 L512,480" stroke="#000" stroke-width="1"/>
+                
+                <path d="M32,0 L32,512" stroke="#000" stroke-width="1"/>
+                <path d="M64,0 L64,512" stroke="#000" stroke-width="1"/>
+                <path d="M96,0 L96,512" stroke="#000" stroke-width="1"/>
+                <path d="M128,0 L128,512" stroke="#000" stroke-width="1"/>
+                <path d="M160,0 L160,512" stroke="#000" stroke-width="1"/>
+                <path d="M192,0 L192,512" stroke="#000" stroke-width="1"/>
+                <path d="M224,0 L224,512" stroke="#000" stroke-width="1"/>
+                <path d="M256,0 L256,512" stroke="#000" stroke-width="1"/>
+                <path d="M288,0 L288,512" stroke="#000" stroke-width="1"/>
+                <path d="M320,0 L320,512" stroke="#000" stroke-width="1"/>
+                <path d="M352,0 L352,512" stroke="#000" stroke-width="1"/>
+                <path d="M384,0 L384,512" stroke="#000" stroke-width="1"/>
+                <path d="M416,0 L416,512" stroke="#000" stroke-width="1"/>
+                <path d="M448,0 L448,512" stroke="#000" stroke-width="1"/>
+                <path d="M480,0 L480,512" stroke="#000" stroke-width="1"/>
+            </g>
         """
         
-        # Add some random paths
-        for i in range(min(num_paths, 20)):  # Limit to 20 paths for the placeholder
-            x = random.randint(50, 462)
-            y = random.randint(50, 462)
-            size = random.randint(20, 100)
-            r = random.randint(0, 255)
-            g = random.randint(0, 255)
-            b = random.randint(0, 255)
-            
-            if random.random() < 0.5:
-                # Circle
-                svg_string += f'<circle cx="{x}" cy="{y}" r="{size/2}" fill="rgb({r},{g},{b})" />'
-            else:
-                # Rectangle
-                svg_string += f'<rect x="{x}" y="{y}" width="{size}" height="{size}" fill="rgb({r},{g},{b})" />'
-        
-        # Add some lines
-        for i in range(min(num_paths // 2, 10)):  # Limit to 10 lines for the placeholder
-            x1 = random.randint(0, 512)
-            y1 = random.randint(0, 512)
-            x2 = random.randint(0, 512)
-            y2 = random.randint(0, 512)
-            r = random.randint(0, 255)
-            g = random.randint(0, 255)
-            b = random.randint(0, 255)
-            width = random.randint(1, 5)
-            
-            svg_string += f'<line x1="{x1}" y1="{y1}" x2="{x2}" y2="{y2}" stroke="rgb({r},{g},{b})" stroke-width="{width}" />'
-        
-        svg_string += "</svg>"
+        # Add some sketch-like paths (DiffSketcher specializes in sketch-like vector graphics)
+        svg_string += '<g filter="url(#pencil-texture)">'
+        
+        # Generate a more complex scene based on the prompt
+        if "car" in prompt.lower():
+            # Generate a car
+            svg_string += self._generate_car_svg(color_ranges)
+        elif "face" in prompt.lower() or "portrait" in prompt.lower():
+            # Generate a face
+            svg_string += self._generate_face_svg(color_ranges)
+        elif "landscape" in prompt.lower() or "mountain" in prompt.lower():
+            # Generate a landscape
+            svg_string += self._generate_landscape_svg(color_ranges)
+        elif "flower" in prompt.lower() or "plant" in prompt.lower():
+            # Generate a flower
+            svg_string += self._generate_flower_svg(color_ranges)
+        elif "animal" in prompt.lower() or "dog" in prompt.lower() or "cat" in prompt.lower():
+            # Generate an animal
+            svg_string += self._generate_animal_svg(color_ranges)
+        else:
+            # Generate abstract art
+            svg_string += self._generate_abstract_svg(color_ranges, num_paths)
+        
+        svg_string += '</g></svg>'
         
         # Convert SVG to PNG
         png_image = self.svg_to_png(svg_string)
         
         return svg_string, png_image
     
+    def _generate_car_svg(self, color_ranges):
+        """Generate a car SVG."""
+        car_svg = ""
+        
+        # Car body
+        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])
+        
+        car_svg += f'<path d="M100,300 Q150,250 200,250 L350,250 Q400,250 450,300 L450,350 Q400,380 350,380 L200,380 Q150,380 100,350 Z" fill="rgb({r},{g},{b})" stroke="black" stroke-width="3" />'
+        
+        # Windows
+        car_svg += '<path d="M150,280 L200,260 L350,260 L400,280 L400,300 L350,320 L200,320 L150,300 Z" fill="#a0d0ff" stroke="black" stroke-width="2" />'
+        
+        # Wheels
+        car_svg += '<circle cx="150" cy="380" r="40" fill="#333" stroke="black" stroke-width="2" />'
+        car_svg += '<circle cx="150" cy="380" r="20" fill="#777" stroke="black" stroke-width="2" />'
+        car_svg += '<circle cx="400" cy="380" r="40" fill="#333" stroke="black" stroke-width="2" />'
+        car_svg += '<circle cx="400" cy="380" r="20" fill="#777" stroke="black" stroke-width="2" />'
+        
+        # Headlights
+        car_svg += '<circle cx="110" cy="320" r="15" fill="#ffff00" stroke="black" stroke-width="2" />'
+        car_svg += '<circle cx="440" cy="320" r="15" fill="#ff0000" stroke="black" stroke-width="2" />'
+        
+        return car_svg
+    
+    def _generate_face_svg(self, color_ranges):
+        """Generate a face SVG."""
+        face_svg = ""
+        
+        # Face shape
+        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])
+        
+        face_svg += f'<ellipse cx="256" cy="256" rx="150" ry="180" fill="rgb({r},{g},{b})" stroke="black" stroke-width="3" />'
+        
+        # Eyes
+        face_svg += '<ellipse cx="200" cy="200" rx="30" ry="20" fill="white" stroke="black" stroke-width="2" />'
+        face_svg += '<circle cx="200" cy="200" r="10" fill="#333" />'
+        face_svg += '<ellipse cx="312" cy="200" rx="30" ry="20" fill="white" stroke="black" stroke-width="2" />'
+        face_svg += '<circle cx="312" cy="200" r="10" fill="#333" />'
+        
+        # Eyebrows
+        face_svg += '<path d="M170,170 Q200,150 230,170" fill="none" stroke="black" stroke-width="3" />'
+        face_svg += '<path d="M282,170 Q312,150 342,170" fill="none" stroke="black" stroke-width="3" />'
+        
+        # Nose
+        face_svg += '<path d="M256,220 Q270,280 256,300 Q242,280 256,220" fill="none" stroke="black" stroke-width="2" />'
+        
+        # Mouth
+        if random.random() < 0.7:  # Smile
+            face_svg += '<path d="M200,320 Q256,380 312,320" fill="none" stroke="black" stroke-width="3" />'
+        else:  # Neutral
+            face_svg += '<path d="M200,330 L312,330" fill="none" stroke="black" stroke-width="3" />'
+        
+        # Hair
+        hair_r = random.randint(0, 100)
+        hair_g = random.randint(0, 100)
+        hair_b = random.randint(0, 100)
+        
+        face_svg += f'<path d="M106,256 Q106,100 256,100 Q406,100 406,256" fill="rgb({hair_r},{hair_g},{hair_b})" stroke="black" stroke-width="3" />'
+        
+        return face_svg
+    
+    def _generate_landscape_svg(self, color_ranges):
+        """Generate a landscape SVG."""
+        landscape_svg = ""
+        
+        # Sky
+        sky_r = random.randint(100, 200)
+        sky_g = random.randint(150, 255)
+        sky_b = random.randint(200, 255)
+        landscape_svg += f'<rect x="0" y="0" width="512" height="300" fill="rgb({sky_r},{sky_g},{sky_b})" />'
+        
+        # Sun
+        sun_x = random.randint(50, 462)
+        sun_y = random.randint(50, 150)
+        landscape_svg += f'<circle cx="{sun_x}" cy="{sun_y}" r="40" fill="#ffff00" />'
+        
+        # Mountains
+        for i in range(5):
+            mountain_x = random.randint(-100, 512)
+            mountain_width = random.randint(200, 400)
+            mountain_height = random.randint(100, 200)
+            
+            r = random.randint(50, 150)
+            g = random.randint(50, 150)
+            b = random.randint(50, 150)
+            
+            landscape_svg += f'<path d="M{mountain_x},{300} L{mountain_x + mountain_width/2},{300 - mountain_height} L{mountain_x + mountain_width},{300} Z" fill="rgb({r},{g},{b})" stroke="black" stroke-width="2" />'
+            
+            # Snow cap
+            landscape_svg += f'<path d="M{mountain_x + mountain_width/4},{300 - mountain_height*0.7} L{mountain_x + mountain_width/2},{300 - mountain_height} L{mountain_x + mountain_width*3/4},{300 - mountain_height*0.7} Z" fill="white" />'
+        
+        # Ground
+        ground_r = random.randint(50, 150)
+        ground_g = random.randint(100, 200)
+        ground_b = random.randint(50, 100)
+        landscape_svg += f'<rect x="0" y="300" width="512" height="212" fill="rgb({ground_r},{ground_g},{ground_b})" />'
+        
+        # Trees
+        for i in range(10):
+            tree_x = random.randint(20, 492)
+            tree_y = random.randint(320, 450)
+            tree_height = random.randint(50, 100)
+            
+            # Trunk
+            landscape_svg += f'<rect x="{tree_x-5}" y="{tree_y}" width="10" height="{tree_height}" fill="#8B4513" />'
+            
+            # Foliage
+            foliage_r = random.randint(0, 100)
+            foliage_g = random.randint(100, 200)
+            foliage_b = random.randint(0, 100)
+            
+            landscape_svg += f'<circle cx="{tree_x}" cy="{tree_y - tree_height/2}" r="{tree_height/2}" fill="rgb({foliage_r},{foliage_g},{foliage_b})" />'
+        
+        return landscape_svg
+    
+    def _generate_flower_svg(self, color_ranges):
+        """Generate a flower SVG."""
+        flower_svg = ""
+        
+        # Stem
+        stem_height = random.randint(150, 300)
+        flower_svg += f'<path d="M256,450 L256,{450-stem_height}" fill="none" stroke="#0a0" stroke-width="5" />'
+        
+        # Leaves
+        leaf_y1 = random.randint(350, 420)
+        leaf_y2 = random.randint(280, 349)
+        
+        flower_svg += f'<path d="M256,{leaf_y1} Q200,{leaf_y1-30} 180,{leaf_y1-10}" fill="none" stroke="#0a0" stroke-width="3" />'
+        flower_svg += f'<path d="M256,{leaf_y2} Q310,{leaf_y2-30} 330,{leaf_y2-10}" fill="none" stroke="#0a0" stroke-width="3" />'
+        
+        # Flower center
+        center_y = 450 - stem_height
+        flower_svg += f'<circle cx="256" cy="{center_y}" r="20" fill="#ff0" stroke="#000" stroke-width="2" />'
+        
+        # Petals
+        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])
+        
+        num_petals = random.randint(5, 12)
+        petal_length = random.randint(40, 70)
+        
+        for i in range(num_petals):
+            angle = 2 * math.pi * i / num_petals
+            petal_x = 256 + petal_length * math.cos(angle)
+            petal_y = center_y + petal_length * math.sin(angle)
+            
+            control_x1 = 256 + petal_length * 0.5 * math.cos(angle - 0.3)
+            control_y1 = center_y + petal_length * 0.5 * math.sin(angle - 0.3)
+            
+            control_x2 = 256 + petal_length * 0.5 * math.cos(angle + 0.3)
+            control_y2 = center_y + petal_length * 0.5 * math.sin(angle + 0.3)
+            
+            flower_svg += f'<path d="M256,{center_y} C{control_x1},{control_y1} {control_x2},{control_y2} {petal_x},{petal_y} C{control_x2},{control_y2} {control_x1},{control_y1} 256,{center_y}" fill="rgb({r},{g},{b})" stroke="#000" stroke-width="1" />'
+        
+        return flower_svg
+    
+    def _generate_animal_svg(self, color_ranges):
+        """Generate an animal SVG."""
+        animal_svg = ""
+        
+        # Body
+        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])
+        
+        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" />'
+        
+        # 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" />'
+        
+        # Nose
+        animal_svg += '<ellipse cx="150" cy="260" rx="15" ry="10" fill="#ffcccc" stroke="black" stroke-width="1" />'
+        
+        # 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" />'
+        
+        # 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" />'
+        
+        # Tail
+        animal_svg += f'<path d="M450,300 Q500,250 520,200" fill="none" stroke="black" stroke-width="3" />'
+        
+        return animal_svg
+    
+    def _generate_abstract_svg(self, color_ranges, num_paths):
+        """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"])
+            
+            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" />'
+            
+            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" />'
+            
+            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" />'
+            
+            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}" />'
+        
+        return abstract_svg
+    
     def __call__(self, data):
         """
         Process the input data and generate SVG output.
@@ -132,7 +447,9 @@ class EndpointHandler:
         negative_prompt = data.get("negative_prompt", "")
         num_paths = int(data.get("num_paths", 96))
         guidance_scale = float(data.get("guidance_scale", 7.5))
-        seed = int(data.get("seed", random.randint(0, 100000)))
+        seed = data.get("seed")
+        if seed is not None:
+            seed = int(seed)
         
         # Generate SVG
         svg_string, png_image = self.generate_svg(