jree423
/

diffsketchedit

@@ -1,295 +0,0 @@
-from typing import Dict, List, Any, Union
-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"DiffSketchEdit 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: List of text prompts for sequential editing
-                - 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, str):
-                inputs = [inputs]
-            elif not isinstance(inputs, list):
-                inputs = ["abstract sketch"]
-            # Extract parameters
-            parameters = data.get("parameters", {})
-            num_paths = parameters.get("num_paths", 96)
-            seed = parameters.get("seed", 1)
-            # Set random seed for reproducibility
-            random.seed(seed)
-            # Generate sequential sketch edits
-            image = self._generate_sequential_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 DiffSketchEdit 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_sequential_sketch(self, prompts: List[str], num_paths: int) -> Image.Image:
-        """Generate a sketch that shows sequential editing based on prompts"""
-        width, height = 224, 224
-        image = Image.new('RGB', (width, height), color='white')
-        draw = ImageDraw.Draw(image)
-        # Process each prompt as a sequential edit
-        for i, prompt in enumerate(prompts[:4]):  # Limit to 4 prompts max
-            alpha = 0.3 + (i * 0.2)  # Increase opacity for later edits
-            self._apply_edit_layer(draw, prompt, width, height, i, alpha, num_paths)
-        return image
-    def _apply_edit_layer(self, draw, prompt: str, width: int, height: int,
-                         layer_index: int, alpha: float, num_paths: int):
-        """Apply an editing layer based on the prompt"""
-        prompt_lower = prompt.lower()
-        # Choose editing operation based on prompt
-        if any(word in prompt_lower for word in ['draw', 'create', 'add']):
-            self._draw_base_sketch(draw, prompt_lower, width, height, layer_index)
-        elif any(word in prompt_lower for word in ['color', 'paint', 'fill']):
-            self._add_color_layer(draw, prompt_lower, width, height, layer_index)
-        elif any(word in prompt_lower for word in ['detail', 'refine', 'enhance']):
-            self._add_detail_layer(draw, prompt_lower, width, height, layer_index)
-        elif any(word in prompt_lower for word in ['modify', 'change', 'edit']):
-            self._modify_existing(draw, prompt_lower, width, height, layer_index)
-        else:
-            self._add_general_elements(draw, prompt_lower, width, height, layer_index)
-    def _draw_base_sketch(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Draw base sketch elements"""
-        if 'cat' in prompt:
-            self._sketch_cat(draw, width, height, layer_index)
-        elif 'house' in prompt:
-            self._sketch_house(draw, width, height, layer_index)
-        elif 'tree' in prompt:
-            self._sketch_tree(draw, width, height, layer_index)
-        elif 'flower' in prompt:
-            self._sketch_flower(draw, width, height, layer_index)
-        else:
-            self._sketch_abstract(draw, width, height, layer_index)
-    def _add_color_layer(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Add color based on prompt"""
-        colors = {
-            'red': '#FF6B6B', 'orange': '#FF8E53', 'yellow': '#FFEAA7',
-            'green': '#55A3FF', 'blue': '#74B9FF', 'purple': '#A29BFE',
-            'pink': '#FD79A8', 'brown': '#FDCB6E'
-        }
-        # Find color in prompt
-        color = '#74B9FF'  # default blue
-        for color_name, color_value in colors.items():
-            if color_name in prompt:
-                color = color_value
-                break
-        # Add color patches
-        for _ in range(3 + layer_index):
-            x = random.randint(20, width-40)
-            y = random.randint(20, height-40)
-            size = random.randint(15, 30)
-            # Create semi-transparent color effect
-            draw.ellipse([x, y, x+size, y+size], fill=color, outline=None)
-    def _add_detail_layer(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Add detail elements"""
-        detail_color = 'black'
-        # Add fine details
-        for _ in range(5 + layer_index * 2):
-            # Random detail lines
-            x1, y1 = random.randint(0, width), random.randint(0, height)
-            x2 = x1 + random.randint(-20, 20)
-            y2 = y1 + random.randint(-20, 20)
-            draw.line([x1, y1, x2, y2], fill=detail_color, width=1)
-        # Add texture dots
-        for _ in range(10 + layer_index * 3):
-            x, y = random.randint(0, width), random.randint(0, height)
-            draw.ellipse([x-1, y-1, x+1, y+1], fill=detail_color)
-    def _modify_existing(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Modify existing elements"""
-        # Add modification strokes
-        for _ in range(3 + layer_index):
-            # Create modification paths
-            start_x = random.randint(width//4, 3*width//4)
-            start_y = random.randint(height//4, 3*height//4)
-            # Draw curved modification
-            points = [start_x, start_y]
-            for step in range(5):
-                start_x += random.randint(-15, 15)
-                start_y += random.randint(-15, 15)
-                start_x = max(10, min(width-10, start_x))
-                start_y = max(10, min(height-10, start_y))
-                points.extend([start_x, start_y])
-            # Draw the modification stroke
-            for i in range(0, len(points)-3, 2):
-                if i+3 < len(points):
-                    draw.line([points[i], points[i+1], points[i+2], points[i+3]],
-                            fill='darkblue', width=2)
-    def _add_general_elements(self, draw, prompt: str, width: int, height: int, layer_index: int):
-        """Add general elements based on prompt"""
-        # Extract key elements from prompt
-        if 'stripe' in prompt:
-            self._add_stripes(draw, width, height, layer_index)
-        elif 'dot' in prompt or 'spot' in prompt:
-            self._add_dots(draw, width, height, layer_index)
-        elif 'line' in prompt:
-            self._add_lines(draw, width, height, layer_index)
-        else:
-            self._add_random_elements(draw, width, height, layer_index)
-    def _sketch_cat(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a cat"""
-        cx, cy = width//2, height//2
-        offset = layer_index * 5  # Slight offset for each layer
-        # Body
-        draw.ellipse([cx-30+offset, cy-5+offset, cx+30+offset, cy+25+offset],
-                    outline='black', width=1)
-        # Head
-        draw.ellipse([cx-20+offset, cy-35+offset, cx+20+offset, cy-5+offset],
-                    outline='black', width=1)
-        # Ears
-        draw.polygon([cx-15+offset, cy-30+offset, cx-8+offset, cy-40+offset, cx-3+offset, cy-30+offset],
-                    outline='black', width=1)
-        draw.polygon([cx+3+offset, cy-30+offset, cx+8+offset, cy-40+offset, cx+15+offset, cy-30+offset],
-                    outline='black', width=1)
-    def _sketch_house(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a house"""
-        offset = layer_index * 3
-        house_x, house_y = width//4 + offset, height//2 + offset
-        house_w, house_h = width//3, height//4
-        # House base
-        draw.rectangle([house_x, house_y, house_x+house_w, house_y+house_h],
-                      outline='black', width=1)
-        # Roof
-        draw.polygon([house_x-5, house_y, house_x+house_w//2, house_y-20, house_x+house_w+5, house_y],
-                    outline='black', width=1)
-    def _sketch_tree(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a tree"""
-        cx, cy = width//2 + layer_index*5, height//2 + layer_index*3
-        # Trunk
-        draw.rectangle([cx-5, cy+10, cx+5, cy+40], outline='black', width=1)
-        # Leaves
-        draw.ellipse([cx-20, cy-15, cx+20, cy+15], outline='black', width=1)
-    def _sketch_flower(self, draw, width: int, height: int, layer_index: int):
-        """Sketch a flower"""
-        cx, cy = width//2 + layer_index*4, height//2 + layer_index*4
-        # Stem
-        draw.line([cx, cy+10, cx, cy+40], fill='black', width=2)
-        # Petals
-        for i in range(6):
-            angle = i * 60
-            x = cx + 15 * math.cos(math.radians(angle))
-            y = cy + 15 * math.sin(math.radians(angle))
-            draw.ellipse([x-5, y-5, x+5, y+5], outline='black', width=1)
-    def _sketch_abstract(self, draw, width: int, height: int, layer_index: int):
-        """Sketch abstract shapes"""
-        for _ in range(3 + layer_index):
-            x, y = random.randint(20, width-20), random.randint(20, height-20)
-            size = random.randint(10, 25)
-            if random.choice([True, False]):
-                draw.ellipse([x, y, x+size, y+size], outline='black', width=1)
-            else:
-                draw.rectangle([x, y, x+size, y+size], outline='black', width=1)
-    def _add_stripes(self, draw, width: int, height: int, layer_index: int):
-        """Add stripe patterns"""
-        stripe_color = ['red', 'blue', 'green', 'orange'][layer_index % 4]
-        spacing = 15 + layer_index * 5
-        for y in range(0, height, spacing):
-            draw.line([0, y, width, y], fill=stripe_color, width=3)
-    def _add_dots(self, draw, width: int, height: int, layer_index: int):
-        """Add dot patterns"""
-        dot_color = ['purple', 'orange', 'green', 'blue'][layer_index % 4]
-        for _ in range(8 + layer_index * 2):
-            x, y = random.randint(10, width-10), random.randint(10, height-10)
-            size = random.randint(3, 8)
-            draw.ellipse([x-size, y-size, x+size, y+size], fill=dot_color)
-    def _add_lines(self, draw, width: int, height: int, layer_index: int):
-        """Add line patterns"""
-        line_color = ['black', 'darkblue', 'darkgreen', 'darkred'][layer_index % 4]
-        for _ in range(5 + layer_index):
-            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=line_color, width=2)
-    def _add_random_elements(self, draw, width: int, height: int, layer_index: int):
-        """Add random elements"""
-        colors = ['red', 'blue', 'green', 'yellow', 'purple', 'orange']
-        for _ in range(4 + layer_index):
-            color = colors[random.randint(0, len(colors)-1)]
-            x, y = random.randint(10, width-20), random.randint(10, height-20)
-            size = random.randint(5, 15)
-            shape = random.choice(['circle', 'square', 'line'])
-            if shape == 'circle':
-                draw.ellipse([x, y, x+size, y+size], fill=color, outline='black')
-            elif shape == 'square':
-                draw.rectangle([x, y, x+size, y+size], fill=color, outline='black')
-            else:
-                draw.line([x, y, x+size, y+size], fill=color, width=3)