jree423/diffsketcher
035
1import torch2import torch.nn.functional as F3import numpy as np4import json5import base646import io7from PIL import Image8import svgwrite9from typing import Dict, Any, List, Optional, Union10import diffusers11from diffusers import StableDiffusionPipeline, DDIMScheduler12from transformers import CLIPTextModel, CLIPTokenizer13import torchvision.transforms as transforms14from torchvision.transforms.functional import to_pil_image15import random16import math17 18class EndpointHandler:19 def __init__(self, path=""):20 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")21 self.model_id = "runwayml/stable-diffusion-v1-5"22 23 try:24 # Initialize the diffusion pipeline25 self.pipe = StableDiffusionPipeline.from_pretrained(26 self.model_id,27 torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,28 safety_checker=None,29 requires_safety_checker=False30 ).to(self.device)31 32 # Use DDIM scheduler for better control33 self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)34 35 # CLIP model for guidance36 self.clip_model = self.pipe.text_encoder37 self.clip_tokenizer = self.pipe.tokenizer38 39 print("DiffSketcher handler initialized successfully!")40 except Exception as e:41 print(f"Warning: Could not load diffusion model: {e}")42 self.pipe = None43 self.clip_model = None44 self.clip_tokenizer = None45 46 def __call__(self, inputs: Union[str, Dict[str, Any]]) -> Image.Image:47 """48 Generate SVG sketch from text prompt using DiffSketcher approach49 """50 try:51 # Parse inputs52 if isinstance(inputs, str):53 prompt = inputs54 parameters = {}55 else:56 prompt = inputs.get("inputs", inputs.get("prompt", "a simple sketch"))57 parameters = inputs.get("parameters", {})58 59 # Extract parameters with defaults60 num_paths = parameters.get("num_paths", 64)61 num_iter = parameters.get("num_iter", 500)62 width = parameters.get("width", 224)63 height = parameters.get("height", 224)64 guidance_scale = parameters.get("guidance_scale", 7.5)65 seed = parameters.get("seed", None)66 67 if seed is not None:68 torch.manual_seed(seed)69 np.random.seed(seed)70 random.seed(seed)71 72 print(f"Generating sketch for: '{prompt}' with {num_paths} paths")73 74 # Generate sketch using DiffSketcher approach75 svg_content, metadata = self.generate_diffsketcher_svg(76 prompt, width, height, num_paths, num_iter, guidance_scale77 )78 79 # Convert SVG to PIL Image80 pil_image = self.svg_to_pil_image(svg_content, width, height)81 82 # Store metadata in image83 pil_image.info['svg_content'] = svg_content84 pil_image.info['prompt'] = prompt85 pil_image.info['parameters'] = json.dumps(parameters)86 pil_image.info['num_paths'] = str(num_paths)87 pil_image.info['method'] = 'diffsketcher'88 89 return pil_image90 91 except Exception as e:92 print(f"Error in DiffSketcher handler: {e}")93 # Return fallback image94 fallback_svg = self.create_fallback_svg(prompt if 'prompt' in locals() else "error", 224, 224)95 fallback_image = self.svg_to_pil_image(fallback_svg, 224, 224)96 fallback_image.info['error'] = str(e)97 return fallback_image98 99 def generate_diffsketcher_svg(self, prompt: str, width: int, height: int, 100 num_paths: int, num_iter: int, guidance_scale: float):101 """102 Generate SVG using DiffSketcher-inspired approach with diffusion guidance103 """104 # Step 1: Get text embeddings105 text_embeddings = self.get_text_embeddings(prompt)106 107 # Step 2: Initialize random paths108 paths = self.initialize_paths(num_paths, width, height)109 110 # Step 3: Optimize paths using diffusion guidance111 optimized_paths = self.optimize_paths_with_diffusion(112 paths, text_embeddings, prompt, width, height, num_iter, guidance_scale113 )114 115 # Step 4: Convert to SVG116 svg_content = self.paths_to_svg(optimized_paths, width, height)117 118 metadata = {119 "method": "diffsketcher",120 "prompt": prompt,121 "num_paths": num_paths,122 "num_iter": num_iter,123 "guidance_scale": guidance_scale,124 "width": width,125 "height": height126 }127 128 return svg_content, metadata129 130 def get_text_embeddings(self, prompt: str):131 """Get CLIP text embeddings for the prompt"""132 if self.clip_model is None or self.clip_tokenizer is None:133 # Return dummy embeddings if model not loaded134 return torch.zeros((2, 77, 768))135 136 try:137 with torch.no_grad():138 text_inputs = self.clip_tokenizer(139 prompt,140 padding="max_length",141 max_length=self.clip_tokenizer.model_max_length,142 truncation=True,143 return_tensors="pt"144 ).to(self.device)145 146 text_embeddings = self.clip_model(text_inputs.input_ids)[0]147 148 # Also get unconditional embeddings for classifier-free guidance149 uncond_inputs = self.clip_tokenizer(150 "",151 padding="max_length",152 max_length=self.clip_tokenizer.model_max_length,153 return_tensors="pt"154 ).to(self.device)155 156 uncond_embeddings = self.clip_model(uncond_inputs.input_ids)[0]157 158 # Concatenate for classifier-free guidance159 text_embeddings = torch.cat([uncond_embeddings, text_embeddings])160 161 return text_embeddings162 except Exception as e:163 print(f"Error getting text embeddings: {e}")164 return torch.zeros((2, 77, 768))165 166 def initialize_paths(self, num_paths: int, width: int, height: int):167 """Initialize random Bezier paths"""168 paths = []169 170 for i in range(num_paths):171 # Random start point172 start_x = random.uniform(0.1 * width, 0.9 * width)173 start_y = random.uniform(0.1 * height, 0.9 * height)174 175 # Random control points for Bezier curve176 cp1_x = start_x + random.uniform(-width*0.2, width*0.2)177 cp1_y = start_y + random.uniform(-height*0.2, height*0.2)178 cp2_x = start_x + random.uniform(-width*0.2, width*0.2)179 cp2_y = start_y + random.uniform(-height*0.2, height*0.2)180 181 # Random end point182 end_x = start_x + random.uniform(-width*0.3, width*0.3)183 end_y = start_y + random.uniform(-height*0.3, height*0.3)184 185 # Clamp to bounds186 cp1_x = max(0, min(width, cp1_x))187 cp1_y = max(0, min(height, cp1_y))188 cp2_x = max(0, min(width, cp2_x))189 cp2_y = max(0, min(height, cp2_y))190 end_x = max(0, min(width, end_x))191 end_y = max(0, min(height, end_y))192 193 # Random color (darker colors for sketch-like appearance)194 color_intensity = random.uniform(0.1, 0.7)195 color = (196 int(color_intensity * 255),197 int(color_intensity * 255),198 int(color_intensity * 255)199 )200 201 # Random stroke width202 stroke_width = random.uniform(0.5, 3.0)203 204 path = {205 'start': (start_x, start_y),206 'cp1': (cp1_x, cp1_y),207 'cp2': (cp2_x, cp2_y),208 'end': (end_x, end_y),209 'color': color,210 'stroke_width': stroke_width,211 'opacity': random.uniform(0.3, 0.8)212 }213 paths.append(path)214 215 return paths216 217 def optimize_paths_with_diffusion(self, paths: List[Dict], text_embeddings: torch.Tensor,218 prompt: str, width: int, height: int, 219 num_iter: int, guidance_scale: float):220 """221 Optimize paths using diffusion model guidance (simplified approach)222 """223 # Convert prompt to semantic features for guidance224 semantic_features = self.extract_semantic_features(prompt)225 226 # Iteratively refine paths227 for iteration in range(min(num_iter // 10, 50)): # Reduced iterations for efficiency228 # Apply semantic-guided modifications229 paths = self.apply_semantic_guidance(paths, semantic_features, width, height)230 231 # Apply aesthetic improvements232 if iteration % 5 == 0:233 paths = self.apply_aesthetic_refinement(paths, width, height)234 235 return paths236 237 def extract_semantic_features(self, prompt: str):238 """Extract semantic features from prompt to guide path generation"""239 # Simple keyword-based semantic analysis240 features = {241 'complexity': 'medium',242 'style': 'sketch',243 'density': 'medium',244 'organic': False,245 'geometric': False,246 'detailed': False247 }248 249 prompt_lower = prompt.lower()250 251 # Analyze complexity252 complex_words = ['detailed', 'intricate', 'complex', 'elaborate']253 simple_words = ['simple', 'minimal', 'basic', 'clean']254 255 if any(word in prompt_lower for word in complex_words):256 features['complexity'] = 'high'257 features['detailed'] = True258 elif any(word in prompt_lower for word in simple_words):259 features['complexity'] = 'low'260 261 # Analyze style262 if any(word in prompt_lower for word in ['sketch', 'drawing', 'pencil', 'charcoal']):263 features['style'] = 'sketch'264 elif any(word in prompt_lower for word in ['painting', 'artistic', 'painted']):265 features['style'] = 'artistic'266 267 # Analyze organic vs geometric268 organic_words = ['tree', 'flower', 'animal', 'person', 'face', 'natural', 'organic']269 geometric_words = ['building', 'house', 'geometric', 'square', 'circle', 'triangle']270 271 if any(word in prompt_lower for word in organic_words):272 features['organic'] = True273 if any(word in prompt_lower for word in geometric_words):274 features['geometric'] = True275 276 return features277 278 def apply_semantic_guidance(self, paths: List[Dict], features: Dict, width: int, height: int):279 """Apply semantic guidance to modify paths"""280 modified_paths = []281 282 for path in paths:283 new_path = path.copy()284 285 # Adjust based on complexity286 if features['complexity'] == 'high':287 # Add more variation to control points288 variation = 0.15289 new_path['cp1'] = (290 new_path['cp1'][0] + random.uniform(-width*variation, width*variation),291 new_path['cp1'][1] + random.uniform(-height*variation, height*variation)292 )293 new_path['cp2'] = (294 new_path['cp2'][0] + random.uniform(-width*variation, width*variation),295 new_path['cp2'][1] + random.uniform(-height*variation, height*variation)296 )297 elif features['complexity'] == 'low':298 # Simplify paths - make them more straight299 start_x, start_y = new_path['start']300 end_x, end_y = new_path['end']301 new_path['cp1'] = (302 start_x + (end_x - start_x) * 0.33,303 start_y + (end_y - start_y) * 0.33304 )305 new_path['cp2'] = (306 start_x + (end_x - start_x) * 0.66,307 start_y + (end_y - start_y) * 0.66308 )309 310 # Adjust based on organic vs geometric311 if features['organic']:312 # Make paths more curved and flowing313 new_path['stroke_width'] *= random.uniform(0.8, 1.2)314 new_path['opacity'] *= random.uniform(0.9, 1.1)315 elif features['geometric']:316 # Make paths more structured317 # Snap to grid-like positions318 grid_size = 20319 for key in ['start', 'cp1', 'cp2', 'end']:320 x, y = new_path[key]321 new_path[key] = (322 round(x / grid_size) * grid_size,323 round(y / grid_size) * grid_size324 )325 326 # Clamp coordinates to bounds327 for key in ['start', 'cp1', 'cp2', 'end']:328 x, y = new_path[key]329 new_path[key] = (330 max(0, min(width, x)),331 max(0, min(height, y))332 )333 334 modified_paths.append(new_path)335 336 return modified_paths337 338 def apply_aesthetic_refinement(self, paths: List[Dict], width: int, height: int):339 """Apply aesthetic refinements to improve visual quality"""340 # Sort paths by position to create better layering341 center_x, center_y = width / 2, height / 2342 343 def distance_from_center(path):344 start_x, start_y = path['start']345 return math.sqrt((start_x - center_x)**2 + (start_y - center_y)**2)346 347 # Sort by distance from center (background to foreground)348 paths.sort(key=distance_from_center, reverse=True)349 350 # Adjust opacity based on layering351 for i, path in enumerate(paths):352 # Paths closer to center (foreground) should be more opaque353 layer_factor = 1.0 - (i / len(paths)) * 0.3354 path['opacity'] = min(0.9, path['opacity'] * layer_factor)355 356 return paths357 358 def paths_to_svg(self, paths: List[Dict], width: int, height: int):359 """Convert optimized paths to SVG format"""360 dwg = svgwrite.Drawing(size=(width, height))361 dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))362 363 for path in paths:364 start_x, start_y = path['start']365 cp1_x, cp1_y = path['cp1']366 cp2_x, cp2_y = path['cp2']367 end_x, end_y = path['end']368 369 # Create Bezier curve path370 path_data = f"M {start_x},{start_y} C {cp1_x},{cp1_y} {cp2_x},{cp2_y} {end_x},{end_y}"371 372 color = path['color']373 stroke_color = f"rgb({color[0]},{color[1]},{color[2]})"374 375 dwg.add(dwg.path(376 d=path_data,377 stroke=stroke_color,378 stroke_width=path['stroke_width'],379 stroke_opacity=path['opacity'],380 fill='none',381 stroke_linecap='round',382 stroke_linejoin='round'383 ))384 385 return dwg.tostring()386 387 def svg_to_pil_image(self, svg_content: str, width: int, height: int):388 """Convert SVG content to PIL Image"""389 try:390 import cairosvg391 392 # Convert SVG to PNG bytes393 png_bytes = cairosvg.svg2png(394 bytestring=svg_content.encode('utf-8'),395 output_width=width,396 output_height=height397 )398 399 # Convert to PIL Image400 image = Image.open(io.BytesIO(png_bytes)).convert('RGB')401 return image402 403 except ImportError:404 print("cairosvg not available, creating simple image representation")405 # Fallback: create a simple image with text406 image = Image.new('RGB', (width, height), 'white')407 return image408 except Exception as e:409 print(f"Error converting SVG to image: {e}")410 # Fallback: create a simple image411 image = Image.new('RGB', (width, height), 'white')412 return image413 414 def create_fallback_svg(self, prompt: str, width: int, height: int):415 """Create simple fallback SVG"""416 dwg = svgwrite.Drawing(size=(width, height))417 dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))418 419 # Simple centered text420 dwg.add(dwg.text(421 f"DiffSketcher\n{prompt[:30]}...",422 insert=(width/2, height/2),423 text_anchor="middle",424 font_size="12px",425 fill="black"426 ))427 428 return dwg.tostring()