import os import sys import torch import base64 import io import json from PIL import Image import svgwrite import numpy as np from diffusers import StableDiffusionPipeline from transformers import CLIPTextModel, CLIPTokenizer import random import math class EndpointHandler: def __init__(self, path=""): """Initialize DiffSketcher handler for Hugging Face Inference API""" self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {self.device}") # Initialize Stable Diffusion pipeline try: self.pipe = StableDiffusionPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 if self.device == "cuda" else torch.float32, safety_checker=None, requires_safety_checker=False ) self.pipe = self.pipe.to(self.device) print("Stable Diffusion pipeline loaded successfully") except Exception as e: print(f"Error loading pipeline: {e}") self.pipe = None # Initialize tokenizer and text encoder try: self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32") self.text_encoder = CLIPTextModel.from_pretrained("openai/clip-vit-base-patch32") self.text_encoder = self.text_encoder.to(self.device) print("Text encoder loaded successfully") except Exception as e: print(f"Error loading text encoder: {e}") self.tokenizer = None self.text_encoder = None def __call__(self, data): """Generate SVG sketch from text prompt""" try: # Extract inputs inputs = data.get("inputs", "") parameters = data.get("parameters", {}) if isinstance(inputs, dict): prompt = inputs.get("prompt", inputs.get("text", "")) else: prompt = str(inputs) if not prompt: prompt = "a simple sketch" # Extract parameters num_paths = parameters.get("num_paths", 96) num_iter = parameters.get("num_iter", 500) guidance_scale = parameters.get("guidance_scale", 7.5) width = parameters.get("width", 224) height = parameters.get("height", 224) seed = parameters.get("seed", 42) # Set seed for reproducibility torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) print(f"Generating SVG for prompt: '{prompt}' with {num_paths} paths") # Generate SVG svg_content = self.generate_svg_sketch( prompt, num_paths, num_iter, guidance_scale, width, height ) # Convert SVG to base64 for transmission svg_base64 = base64.b64encode(svg_content.encode('utf-8')).decode('utf-8') # Return result result = { "svg": svg_content, "svg_base64": svg_base64, "prompt": prompt, "parameters": { "num_paths": num_paths, "num_iter": num_iter, "guidance_scale": guidance_scale, "width": width, "height": height, "seed": seed } } return result except Exception as e: print(f"Error in handler: {e}") # Return a simple fallback SVG fallback_svg = self.create_fallback_svg(prompt, width, height) return { "svg": fallback_svg, "svg_base64": base64.b64encode(fallback_svg.encode('utf-8')).decode('utf-8'), "prompt": prompt, "error": str(e) } def generate_svg_sketch(self, prompt, num_paths, num_iter, guidance_scale, width, height): """Generate SVG sketch using simplified DiffSketcher approach""" try: # Get text embeddings text_embeddings = self.get_text_embeddings(prompt) # Generate attention maps (simplified) attention_maps = self.generate_attention_maps(prompt, width, height) # Initialize SVG paths based on attention paths = self.initialize_paths_from_attention(attention_maps, num_paths, width, height) # Optimize paths (simplified version) optimized_paths = self.optimize_paths(paths, text_embeddings, num_iter, guidance_scale) # Create SVG svg_content = self.create_svg_from_paths(optimized_paths, width, height) return svg_content except Exception as e: print(f"Error in generate_svg_sketch: {e}") return self.create_fallback_svg(prompt, width, height) def get_text_embeddings(self, prompt): """Get text embeddings from CLIP""" if self.tokenizer is None or self.text_encoder is None: return None try: inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True) inputs = {k: v.to(self.device) for k, v in inputs.items()} with torch.no_grad(): embeddings = self.text_encoder(**inputs).last_hidden_state return embeddings except Exception as e: print(f"Error getting text embeddings: {e}") return None def generate_attention_maps(self, prompt, width, height): """Generate simplified attention maps""" # Create attention maps based on prompt keywords attention_map = np.zeros((height, width)) # Simple keyword-based attention keywords = prompt.lower().split() for i, keyword in enumerate(keywords[:5]): # Limit to 5 keywords # Create attention region for each keyword center_x = (i + 1) * width // (len(keywords) + 1) center_y = height // 2 # Create Gaussian-like attention y, x = np.ogrid[:height, :width] mask = ((x - center_x) ** 2 + (y - center_y) ** 2) < (min(width, height) // 4) ** 2 attention_map[mask] += 1.0 # Normalize if attention_map.max() > 0: attention_map = attention_map / attention_map.max() return attention_map def initialize_paths_from_attention(self, attention_map, num_paths, width, height): """Initialize SVG paths based on attention maps""" paths = [] # Find high attention regions threshold = 0.3 high_attention = attention_map > threshold if not np.any(high_attention): # Fallback: create random paths return self.create_random_paths(num_paths, width, height) # Get coordinates of high attention regions y_coords, x_coords = np.where(high_attention) for i in range(num_paths): if len(x_coords) > 0: # Sample random points from high attention regions idx = np.random.choice(len(x_coords), size=min(4, len(x_coords)), replace=False) path_points = [(x_coords[j], y_coords[j]) for j in idx] # Sort points to create a reasonable path path_points.sort(key=lambda p: p[0]) paths.append(path_points) else: # Fallback to random path paths.append(self.create_single_random_path(width, height)) return paths def create_random_paths(self, num_paths, width, height): """Create random paths as fallback""" paths = [] for i in range(num_paths): paths.append(self.create_single_random_path(width, height)) return paths def create_single_random_path(self, width, height): """Create a single random path""" num_points = random.randint(3, 6) points = [] for _ in range(num_points): x = random.randint(0, width) y = random.randint(0, height) points.append((x, y)) return points def optimize_paths(self, paths, text_embeddings, num_iter, guidance_scale): """Simplified path optimization""" # For now, just add some smoothing and variation optimized_paths = [] for path in paths: if len(path) < 2: optimized_paths.append(path) continue # Add some smoothing smoothed_path = [] for i in range(len(path)): if i == 0 or i == len(path) - 1: smoothed_path.append(path[i]) else: # Simple smoothing prev_x, prev_y = path[i-1] curr_x, curr_y = path[i] next_x, next_y = path[i+1] smooth_x = (prev_x + curr_x + next_x) / 3 smooth_y = (prev_y + curr_y + next_y) / 3 smoothed_path.append((smooth_x, smooth_y)) optimized_paths.append(smoothed_path) return optimized_paths def create_svg_from_paths(self, paths, width, height): """Create SVG content from optimized paths""" dwg = svgwrite.Drawing(size=(width, height)) # Add white background dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white')) # Add paths for i, path in enumerate(paths): if len(path) < 2: continue # Create path string path_str = f"M {path[0][0]},{path[0][1]}" for point in path[1:]: path_str += f" L {point[0]},{point[1]}" # Vary stroke properties stroke_width = random.uniform(0.5, 3.0) stroke_color = f"rgb({random.randint(0, 100)},{random.randint(0, 100)},{random.randint(0, 100)})" dwg.add(dwg.path( d=path_str, stroke=stroke_color, stroke_width=stroke_width, fill='none', stroke_linecap='round', stroke_linejoin='round' )) return dwg.tostring() def create_fallback_svg(self, prompt, width=224, height=224): """Create a simple fallback SVG""" dwg = svgwrite.Drawing(size=(width, height)) # Add white background dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white')) # Add simple sketch based on prompt prompt_lower = prompt.lower() if any(word in prompt_lower for word in ['mountain', 'landscape']): self._add_mountain_sketch(dwg, width, height) elif any(word in prompt_lower for word in ['house', 'building']): self._add_house_sketch(dwg, width, height) elif any(word in prompt_lower for word in ['flower', 'plant']): self._add_flower_sketch(dwg, width, height) else: self._add_abstract_sketch(dwg, width, height, prompt) return dwg.tostring() def _add_mountain_sketch(self, dwg, width, height): """Add mountain sketch to SVG""" # Mountain outline points = [(0, height*0.7)] for x in range(0, width, 20): y = height * 0.7 + 30 * math.sin(x * 0.02) + 15 * math.sin(x * 0.05) points.append((x, y)) points.append((width, height)) points.append((0, height)) dwg.add(dwg.polygon(points, fill='lightgray', stroke='black', stroke_width=2)) def _add_house_sketch(self, dwg, width, height): """Add house sketch to SVG""" # House base house_width = width * 0.6 house_height = height * 0.4 house_x = (width - house_width) / 2 house_y = height * 0.4 dwg.add(dwg.rect( insert=(house_x, house_y), size=(house_width, house_height), fill='lightblue', stroke='black', stroke_width=2 )) # Roof roof_points = [ (house_x, house_y), (house_x + house_width/2, house_y - house_height*0.3), (house_x + house_width, house_y) ] dwg.add(dwg.polygon(roof_points, fill='red', stroke='black', stroke_width=2)) def _add_flower_sketch(self, dwg, width, height): """Add flower sketch to SVG""" center_x, center_y = width/2, height/2 # Stem dwg.add(dwg.line( start=(center_x, center_y + 20), end=(center_x, height - 20), stroke='green', stroke_width=4 )) # Petals for angle in range(0, 360, 45): x = center_x + 25 * math.cos(math.radians(angle)) y = center_y + 25 * math.sin(math.radians(angle)) dwg.add(dwg.circle( center=(x, y), r=8, fill='pink', stroke='red', stroke_width=1 )) # Center dwg.add(dwg.circle( center=(center_x, center_y), r=8, fill='yellow', stroke='orange', stroke_width=2 )) def _add_abstract_sketch(self, dwg, width, height, prompt): """Add abstract sketch to SVG""" # Create flowing lines based on prompt hash prompt_hash = hash(prompt) % 100 for i in range(8): points = [] start_x = (i * 30 + prompt_hash) % (width - 40) + 20 start_y = (i * 25 + prompt_hash) % (height - 40) + 20 for j in range(4): x = start_x + j * 25 + 15 * math.sin((i + j + prompt_hash) * 0.5) y = start_y + j * 20 + 15 * math.cos((i + j + prompt_hash) * 0.3) points.append((max(0, min(width, x)), max(0, min(height, y)))) # Create path if len(points) > 1: path_str = f"M {points[0][0]},{points[0][1]}" for point in points[1:]: path_str += f" L {point[0]},{point[1]}" color_val = (i * 30) % 200 + 50 dwg.add(dwg.path( d=path_str, stroke=f"rgb({color_val},{color_val//2},{color_val//3})", stroke_width=2, fill='none', stroke_linecap='round' ))