import torch import torch.nn.functional as F import numpy as np import json import base64 import io from PIL import Image import svgwrite from typing import Dict, Any, List, Optional, Union import diffusers from diffusers import StableDiffusionPipeline, DDIMScheduler from transformers import CLIPTextModel, CLIPTokenizer import torchvision.transforms as transforms from torchvision.transforms.functional import to_pil_image import random import math class EndpointHandler: def __init__(self, path=""): self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.model_id = "runwayml/stable-diffusion-v1-5" try: # Initialize the diffusion pipeline self.pipe = StableDiffusionPipeline.from_pretrained( self.model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, safety_checker=None, requires_safety_checker=False ).to(self.device) # Use DDIM scheduler for better control self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config) # CLIP model for guidance self.clip_model = self.pipe.text_encoder self.clip_tokenizer = self.pipe.tokenizer print("DiffSketcher handler initialized successfully!") except Exception as e: print(f"Warning: Could not load diffusion model: {e}") self.pipe = None self.clip_model = None self.clip_tokenizer = None def __call__(self, inputs: Union[str, Dict[str, Any]]) -> Image.Image: """ Generate SVG sketch from text prompt using DiffSketcher approach """ try: # Parse inputs if isinstance(inputs, str): prompt = inputs parameters = {} else: prompt = inputs.get("inputs", inputs.get("prompt", "a simple sketch")) parameters = inputs.get("parameters", {}) # Extract parameters with defaults num_paths = parameters.get("num_paths", 64) num_iter = parameters.get("num_iter", 500) width = parameters.get("width", 224) height = parameters.get("height", 224) guidance_scale = parameters.get("guidance_scale", 7.5) seed = parameters.get("seed", None) if seed is not None: torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) print(f"Generating sketch for: '{prompt}' with {num_paths} paths") # Generate sketch using DiffSketcher approach svg_content, metadata = self.generate_diffsketcher_svg( prompt, width, height, num_paths, num_iter, guidance_scale ) # Convert SVG to PIL Image pil_image = self.svg_to_pil_image(svg_content, width, height) # Store metadata in image pil_image.info['svg_content'] = svg_content pil_image.info['prompt'] = prompt pil_image.info['parameters'] = json.dumps(parameters) pil_image.info['num_paths'] = str(num_paths) pil_image.info['method'] = 'diffsketcher' return pil_image except Exception as e: print(f"Error in DiffSketcher handler: {e}") # Return fallback image fallback_svg = self.create_fallback_svg(prompt if 'prompt' in locals() else "error", 224, 224) fallback_image = self.svg_to_pil_image(fallback_svg, 224, 224) fallback_image.info['error'] = str(e) return fallback_image def generate_diffsketcher_svg(self, prompt: str, width: int, height: int, num_paths: int, num_iter: int, guidance_scale: float): """ Generate SVG using DiffSketcher-inspired approach with diffusion guidance """ # Step 1: Get text embeddings text_embeddings = self.get_text_embeddings(prompt) # Step 2: Initialize random paths paths = self.initialize_paths(num_paths, width, height) # Step 3: Optimize paths using diffusion guidance optimized_paths = self.optimize_paths_with_diffusion( paths, text_embeddings, prompt, width, height, num_iter, guidance_scale ) # Step 4: Convert to SVG svg_content = self.paths_to_svg(optimized_paths, width, height) metadata = { "method": "diffsketcher", "prompt": prompt, "num_paths": num_paths, "num_iter": num_iter, "guidance_scale": guidance_scale, "width": width, "height": height } return svg_content, metadata def get_text_embeddings(self, prompt: str): """Get CLIP text embeddings for the prompt""" if self.clip_model is None or self.clip_tokenizer is None: # Return dummy embeddings if model not loaded return torch.zeros((2, 77, 768)) try: with torch.no_grad(): text_inputs = self.clip_tokenizer( prompt, padding="max_length", max_length=self.clip_tokenizer.model_max_length, truncation=True, return_tensors="pt" ).to(self.device) text_embeddings = self.clip_model(text_inputs.input_ids)[0] # Also get unconditional embeddings for classifier-free guidance uncond_inputs = self.clip_tokenizer( "", padding="max_length", max_length=self.clip_tokenizer.model_max_length, return_tensors="pt" ).to(self.device) uncond_embeddings = self.clip_model(uncond_inputs.input_ids)[0] # Concatenate for classifier-free guidance text_embeddings = torch.cat([uncond_embeddings, text_embeddings]) return text_embeddings except Exception as e: print(f"Error getting text embeddings: {e}") return torch.zeros((2, 77, 768)) def initialize_paths(self, num_paths: int, width: int, height: int): """Initialize random Bezier paths""" paths = [] for i in range(num_paths): # Random start point start_x = random.uniform(0.1 * width, 0.9 * width) start_y = random.uniform(0.1 * height, 0.9 * height) # Random control points for Bezier curve cp1_x = start_x + random.uniform(-width*0.2, width*0.2) cp1_y = start_y + random.uniform(-height*0.2, height*0.2) cp2_x = start_x + random.uniform(-width*0.2, width*0.2) cp2_y = start_y + random.uniform(-height*0.2, height*0.2) # Random end point end_x = start_x + random.uniform(-width*0.3, width*0.3) end_y = start_y + random.uniform(-height*0.3, height*0.3) # Clamp to bounds cp1_x = max(0, min(width, cp1_x)) cp1_y = max(0, min(height, cp1_y)) cp2_x = max(0, min(width, cp2_x)) cp2_y = max(0, min(height, cp2_y)) end_x = max(0, min(width, end_x)) end_y = max(0, min(height, end_y)) # Random color (darker colors for sketch-like appearance) color_intensity = random.uniform(0.1, 0.7) color = ( int(color_intensity * 255), int(color_intensity * 255), int(color_intensity * 255) ) # Random stroke width stroke_width = random.uniform(0.5, 3.0) path = { 'start': (start_x, start_y), 'cp1': (cp1_x, cp1_y), 'cp2': (cp2_x, cp2_y), 'end': (end_x, end_y), 'color': color, 'stroke_width': stroke_width, 'opacity': random.uniform(0.3, 0.8) } paths.append(path) return paths def optimize_paths_with_diffusion(self, paths: List[Dict], text_embeddings: torch.Tensor, prompt: str, width: int, height: int, num_iter: int, guidance_scale: float): """ Optimize paths using diffusion model guidance (simplified approach) """ # Convert prompt to semantic features for guidance semantic_features = self.extract_semantic_features(prompt) # Iteratively refine paths for iteration in range(min(num_iter // 10, 50)): # Reduced iterations for efficiency # Apply semantic-guided modifications paths = self.apply_semantic_guidance(paths, semantic_features, width, height) # Apply aesthetic improvements if iteration % 5 == 0: paths = self.apply_aesthetic_refinement(paths, width, height) return paths def extract_semantic_features(self, prompt: str): """Extract semantic features from prompt to guide path generation""" # Simple keyword-based semantic analysis features = { 'complexity': 'medium', 'style': 'sketch', 'density': 'medium', 'organic': False, 'geometric': False, 'detailed': False } prompt_lower = prompt.lower() # Analyze complexity complex_words = ['detailed', 'intricate', 'complex', 'elaborate'] simple_words = ['simple', 'minimal', 'basic', 'clean'] if any(word in prompt_lower for word in complex_words): features['complexity'] = 'high' features['detailed'] = True elif any(word in prompt_lower for word in simple_words): features['complexity'] = 'low' # Analyze style if any(word in prompt_lower for word in ['sketch', 'drawing', 'pencil', 'charcoal']): features['style'] = 'sketch' elif any(word in prompt_lower for word in ['painting', 'artistic', 'painted']): features['style'] = 'artistic' # Analyze organic vs geometric organic_words = ['tree', 'flower', 'animal', 'person', 'face', 'natural', 'organic'] geometric_words = ['building', 'house', 'geometric', 'square', 'circle', 'triangle'] if any(word in prompt_lower for word in organic_words): features['organic'] = True if any(word in prompt_lower for word in geometric_words): features['geometric'] = True return features def apply_semantic_guidance(self, paths: List[Dict], features: Dict, width: int, height: int): """Apply semantic guidance to modify paths""" modified_paths = [] for path in paths: new_path = path.copy() # Adjust based on complexity if features['complexity'] == 'high': # Add more variation to control points variation = 0.15 new_path['cp1'] = ( new_path['cp1'][0] + random.uniform(-width*variation, width*variation), new_path['cp1'][1] + random.uniform(-height*variation, height*variation) ) new_path['cp2'] = ( new_path['cp2'][0] + random.uniform(-width*variation, width*variation), new_path['cp2'][1] + random.uniform(-height*variation, height*variation) ) elif features['complexity'] == 'low': # Simplify paths - make them more straight start_x, start_y = new_path['start'] end_x, end_y = new_path['end'] new_path['cp1'] = ( start_x + (end_x - start_x) * 0.33, start_y + (end_y - start_y) * 0.33 ) new_path['cp2'] = ( start_x + (end_x - start_x) * 0.66, start_y + (end_y - start_y) * 0.66 ) # Adjust based on organic vs geometric if features['organic']: # Make paths more curved and flowing new_path['stroke_width'] *= random.uniform(0.8, 1.2) new_path['opacity'] *= random.uniform(0.9, 1.1) elif features['geometric']: # Make paths more structured # Snap to grid-like positions grid_size = 20 for key in ['start', 'cp1', 'cp2', 'end']: x, y = new_path[key] new_path[key] = ( round(x / grid_size) * grid_size, round(y / grid_size) * grid_size ) # Clamp coordinates to bounds for key in ['start', 'cp1', 'cp2', 'end']: x, y = new_path[key] new_path[key] = ( max(0, min(width, x)), max(0, min(height, y)) ) modified_paths.append(new_path) return modified_paths def apply_aesthetic_refinement(self, paths: List[Dict], width: int, height: int): """Apply aesthetic refinements to improve visual quality""" # Sort paths by position to create better layering center_x, center_y = width / 2, height / 2 def distance_from_center(path): start_x, start_y = path['start'] return math.sqrt((start_x - center_x)**2 + (start_y - center_y)**2) # Sort by distance from center (background to foreground) paths.sort(key=distance_from_center, reverse=True) # Adjust opacity based on layering for i, path in enumerate(paths): # Paths closer to center (foreground) should be more opaque layer_factor = 1.0 - (i / len(paths)) * 0.3 path['opacity'] = min(0.9, path['opacity'] * layer_factor) return paths def paths_to_svg(self, paths: List[Dict], width: int, height: int): """Convert optimized paths to SVG format""" dwg = svgwrite.Drawing(size=(width, height)) dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white')) for path in paths: start_x, start_y = path['start'] cp1_x, cp1_y = path['cp1'] cp2_x, cp2_y = path['cp2'] end_x, end_y = path['end'] # Create Bezier curve path path_data = f"M {start_x},{start_y} C {cp1_x},{cp1_y} {cp2_x},{cp2_y} {end_x},{end_y}" color = path['color'] stroke_color = f"rgb({color[0]},{color[1]},{color[2]})" dwg.add(dwg.path( d=path_data, stroke=stroke_color, stroke_width=path['stroke_width'], stroke_opacity=path['opacity'], fill='none', stroke_linecap='round', stroke_linejoin='round' )) return dwg.tostring() def svg_to_pil_image(self, svg_content: str, width: int, height: int): """Convert SVG content to PIL Image""" try: import cairosvg # Convert SVG to PNG bytes png_bytes = cairosvg.svg2png( bytestring=svg_content.encode('utf-8'), output_width=width, output_height=height ) # Convert to PIL Image image = Image.open(io.BytesIO(png_bytes)).convert('RGB') return image except ImportError: print("cairosvg not available, creating simple image representation") # Fallback: create a simple image with text image = Image.new('RGB', (width, height), 'white') return image except Exception as e: print(f"Error converting SVG to image: {e}") # Fallback: create a simple image image = Image.new('RGB', (width, height), 'white') return image def create_fallback_svg(self, prompt: str, width: int, height: int): """Create simple fallback SVG""" dwg = svgwrite.Drawing(size=(width, height)) dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white')) # Simple centered text dwg.add(dwg.text( f"DiffSketcher\n{prompt[:30]}...", insert=(width/2, height/2), text_anchor="middle", font_size="12px", fill="black" )) return dwg.tostring()