jree423
/

diffsketcher

@@ -1,238 +0,0 @@
-from typing import Dict, List, Any
-import torch
-from PIL import Image, ImageDraw
-import io
-import base64
-import random
-import math
-class EndpointHandler:
-    def __init__(self, path=""):
-        """Initialize the handler with minimal dependencies"""
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        print(f"DiffSketcher handler initialized on {self.device}")
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        """
-        Process the request and return generated image
-        Args:
-            data: Dictionary containing:
-                - inputs: Text prompt for sketch generation
-                - parameters: Optional parameters (num_paths, num_iter, etc.)
-        Returns:
-            List containing dictionary with base64 encoded image
-        """
-        try:
-            # Extract inputs
-            inputs = data.get("inputs", "")
-            if isinstance(inputs, list):
-                inputs = inputs[0] if inputs else ""
-            # Extract parameters
-            parameters = data.get("parameters", {})
-            num_paths = parameters.get("num_paths", 64)
-            seed = parameters.get("seed", 42)
-            # Set random seed for reproducibility
-            random.seed(seed)
-            # Generate SVG-style sketch
-            image = self._generate_sketch(inputs, num_paths)
-            # Convert to base64
-            buffered = io.BytesIO()
-            image.save(buffered, format="PNG")
-            img_base64 = base64.b64encode(buffered.getvalue()).decode()
-            return [{"generated_image": img_base64}]
-        except Exception as e:
-            print(f"Error in DiffSketcher handler: {e}")
-            # Return error image
-            error_img = Image.new('RGB', (224, 224), color='lightcoral')
-            draw = ImageDraw.Draw(error_img)
-            draw.text((10, 100), f"Error: {str(e)[:30]}", fill='white')
-            buffered = io.BytesIO()
-            error_img.save(buffered, format="PNG")
-            img_base64 = base64.b64encode(buffered.getvalue()).decode()
-            return [{"generated_image": img_base64}]
-    def _generate_sketch(self, prompt: str, num_paths: int) -> Image.Image:
-        """Generate a sketch-style image based on the prompt"""
-        # Create canvas
-        width, height = 224, 224
-        image = Image.new('RGB', (width, height), color='white')
-        draw = ImageDraw.Draw(image)
-        # Simple prompt-based sketch generation
-        prompt_lower = prompt.lower()
-        # Generate sketch elements based on prompt keywords
-        if any(word in prompt_lower for word in ['mountain', 'landscape', 'hill']):
-            self._draw_mountains(draw, width, height, num_paths)
-        elif any(word in prompt_lower for word in ['cat', 'animal', 'pet']):
-            self._draw_cat(draw, width, height, num_paths)
-        elif any(word in prompt_lower for word in ['flower', 'plant', 'garden']):
-            self._draw_flower(draw, width, height, num_paths)
-        elif any(word in prompt_lower for word in ['house', 'building', 'home']):
-            self._draw_house(draw, width, height, num_paths)
-        elif any(word in prompt_lower for word in ['tree', 'forest', 'wood']):
-            self._draw_tree(draw, width, height, num_paths)
-        else:
-            self._draw_abstract(draw, width, height, num_paths)
-        return image
-    def _draw_mountains(self, draw, width, height, num_paths):
-        """Draw mountain landscape"""
-        # Background mountains
-        for i in range(3):
-            y_offset = height // 3 + i * 20
-            points = []
-            for x in range(0, width + 20, 20):
-                y = y_offset + random.randint(-30, 10)
-                points.extend([x, y])
-            if len(points) >= 6:
-                draw.polygon(points + [width, height, 0, height],
-                           fill=f'rgb({200-i*30},{220-i*20},{240-i*10})',
-                           outline='gray')
-        # Add some trees
-        for _ in range(num_paths // 20):
-            x = random.randint(10, width-10)
-            y = random.randint(height//2, height-20)
-            draw.ellipse([x-5, y-15, x+5, y], fill='darkgreen')
-            draw.line([x, y, x, y+15], fill='brown', width=2)
-    def _draw_cat(self, draw, width, height, num_paths):
-        """Draw a simple cat"""
-        cx, cy = width//2, height//2
-        # Body
-        draw.ellipse([cx-40, cy-10, cx+40, cy+30], outline='black', width=2)
-        # Head
-        draw.ellipse([cx-25, cy-40, cx+25, cy-10], outline='black', width=2)
-        # Ears
-        draw.polygon([cx-20, cy-35, cx-10, cy-50, cx-5, cy-35], outline='black', width=2)
-        draw.polygon([cx+5, cy-35, cx+10, cy-50, cx+20, cy-35], outline='black', width=2)
-        # Eyes
-        draw.ellipse([cx-15, cy-30, cx-10, cy-25], fill='black')
-        draw.ellipse([cx+10, cy-30, cx+15, cy-25], fill='black')
-        # Nose
-        draw.polygon([cx-2, cy-20, cx+2, cy-20, cx, cy-15], fill='pink')
-        # Whiskers
-        for i in range(3):
-            y_pos = cy-18 + i*3
-            draw.line([cx-25, y_pos, cx-35, y_pos], fill='black', width=1)
-            draw.line([cx+25, y_pos, cx+35, y_pos], fill='black', width=1)
-        # Tail
-        draw.arc([cx+30, cy-5, cx+60, cy+25], 0, 180, fill='black', width=3)
-    def _draw_flower(self, draw, width, height, num_paths):
-        """Draw a simple flower"""
-        cx, cy = width//2, height//2
-        # Stem
-        draw.line([cx, cy+20, cx, height-20], fill='green', width=4)
-        # Petals
-        petal_colors = ['red', 'pink', 'yellow', 'orange', 'purple']
-        for i in range(6):
-            angle = i * 60
-            x1 = cx + 20 * math.cos(math.radians(angle))
-            y1 = cy + 20 * math.sin(math.radians(angle))
-            x2 = cx + 35 * math.cos(math.radians(angle))
-            y2 = cy + 35 * math.sin(math.radians(angle))
-            color = random.choice(petal_colors)
-            draw.ellipse([x2-8, y2-8, x2+8, y2+8], fill=color, outline='darkred')
-        # Center
-        draw.ellipse([cx-8, cy-8, cx+8, cy+8], fill='yellow', outline='orange')
-        # Leaves
-        draw.ellipse([cx-15, cy+30, cx-5, cy+50], fill='green', outline='darkgreen')
-        draw.ellipse([cx+5, cy+35, cx+15, cy+55], fill='green', outline='darkgreen')
-    def _draw_house(self, draw, width, height, num_paths):
-        """Draw a simple house"""
-        # House base
-        house_x, house_y = width//4, height//2
-        house_w, house_h = width//2, height//3
-        draw.rectangle([house_x, house_y, house_x+house_w, house_y+house_h],
-                      outline='black', width=2)
-        # Roof
-        draw.polygon([house_x-10, house_y, width//2, house_y-40, house_x+house_w+10, house_y],
-                    fill='red', outline='darkred', width=2)
-        # Door
-        door_w, door_h = 20, 40
-        door_x = house_x + house_w//2 - door_w//2
-        door_y = house_y + house_h - door_h
-        draw.rectangle([door_x, door_y, door_x+door_w, door_y+door_h],
-                      fill='brown', outline='black', width=2)
-        # Windows
-        win_size = 15
-        draw.rectangle([house_x+10, house_y+15, house_x+10+win_size, house_y+15+win_size],
-                      fill='lightblue', outline='black', width=2)
-        draw.rectangle([house_x+house_w-25, house_y+15, house_x+house_w-10, house_y+30],
-                      fill='lightblue', outline='black', width=2)
-        # Chimney
-        draw.rectangle([house_x+house_w-20, house_y-35, house_x+house_w-10, house_y-10],
-                      fill='gray', outline='black', width=2)
-    def _draw_tree(self, draw, width, height, num_paths):
-        """Draw a simple tree"""
-        cx, cy = width//2, height//2
-        # Trunk
-        trunk_w, trunk_h = 15, 60
-        draw.rectangle([cx-trunk_w//2, cy+20, cx+trunk_w//2, cy+20+trunk_h],
-                      fill='brown', outline='black', width=2)
-        # Leaves (multiple circles for fuller look)
-        leaf_positions = [
-            (cx, cy-20), (cx-25, cy-10), (cx+25, cy-10),
-            (cx-15, cy+5), (cx+15, cy+5), (cx, cy+15)
-        ]
-        for lx, ly in leaf_positions:
-            size = random.randint(20, 35)
-            draw.ellipse([lx-size//2, ly-size//2, lx+size//2, ly+size//2],
-                        fill='green', outline='darkgreen', width=1)
-    def _draw_abstract(self, draw, width, height, num_paths):
-        """Draw abstract shapes"""
-        colors = ['red', 'blue', 'green', 'yellow', 'purple', 'orange', 'pink']
-        # Draw random shapes
-        for _ in range(min(num_paths//10, 8)):
-            color = random.choice(colors)
-            shape_type = random.choice(['circle', 'rectangle', 'line'])
-            if shape_type == 'circle':
-                x, y = random.randint(20, width-20), random.randint(20, height-20)
-                r = random.randint(10, 30)
-                draw.ellipse([x-r, y-r, x+r, y+r], fill=color, outline='black')
-            elif shape_type == 'rectangle':
-                x1, y1 = random.randint(10, width//2), random.randint(10, height//2)
-                x2, y2 = random.randint(width//2, width-10), random.randint(height//2, height-10)
-                draw.rectangle([x1, y1, x2, y2], fill=color, outline='black')
-            else:  # line
-                x1, y1 = random.randint(0, width), random.randint(0, height)
-                x2, y2 = random.randint(0, width), random.randint(0, height)
-                draw.line([x1, y1, x2, y2], fill=color, width=random.randint(2, 5))