File size: 30,706 Bytes
b0efdb8
f9319c0
b0efdb8
f9319c0
 
 
 
b0efdb8
f9319c0
 
 
 
 
26a4cfd
b0efdb8
f9319c0
26a4cfd
b6bfd0d
 
f9319c0
 
 
b6bfd0d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26a4cfd
f9319c0
 
 
 
b0efdb8
f9319c0
 
e124116
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f9319c0
 
 
 
e124116
f9319c0
 
e124116
 
 
f9319c0
 
b0efdb8
f9319c0
b0efdb8
 
f9319c0
 
b0efdb8
f9319c0
 
 
 
b0efdb8
 
 
 
 
ad55c44
b0efdb8
ad55c44
b0efdb8
ad55c44
b0efdb8
ad55c44
b0efdb8
 
ad55c44
b0efdb8
ad55c44
 
 
 
f9319c0
ad55c44
 
 
 
 
 
 
b0efdb8
 
 
ad55c44
b0efdb8
ad55c44
 
 
 
b0efdb8
 
 
 
 
 
f9319c0
 
b0efdb8
 
f9319c0
b0efdb8
 
 
 
 
 
 
 
ad55c44
b0efdb8
 
 
 
 
 
 
ad55c44
 
b0efdb8
 
ad55c44
 
 
b0efdb8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ad55c44
b0efdb8
 
 
 
ad55c44
 
b0efdb8
 
ad55c44
 
 
b0efdb8
 
 
 
 
 
f9319c0
b0efdb8
f9319c0
b0efdb8
 
 
 
 
 
 
 
 
 
ad55c44
b0efdb8
 
 
 
 
 
ad55c44
 
b0efdb8
 
ad55c44
 
 
b0efdb8
 
 
 
 
 
 
 
ad55c44
b0efdb8
 
 
 
ad55c44
 
b0efdb8
 
ad55c44
 
 
b0efdb8
 
 
 
 
 
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
26a4cfd
f9319c0
b0efdb8
 
 
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
 
 
 
 
 
b0efdb8
 
26a4cfd
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
 
26a4cfd
b0efdb8
 
f9319c0
 
 
b0efdb8
 
26a4cfd
f9319c0
 
 
26a4cfd
f9319c0
26a4cfd
b0efdb8
26a4cfd
b0efdb8
f9319c0
b0efdb8
 
26a4cfd
f9319c0
b0efdb8
 
 
f9319c0
b0efdb8
 
f9319c0
26a4cfd
f9319c0
 
b0efdb8
 
 
f9319c0
 
 
b0efdb8
f9319c0
 
 
26a4cfd
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
26a4cfd
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
 
b0efdb8
 
 
 
26a4cfd
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
26a4cfd
f9319c0
 
 
 
b0efdb8
 
 
f9319c0
b0efdb8
 
 
26a4cfd
f9319c0
 
 
 
b0efdb8
26a4cfd
f9319c0
 
 
 
 
b0efdb8
f9319c0
b0efdb8
f9319c0
b0efdb8
f9319c0
b0efdb8
 
f9319c0
 
 
 
 
 
 
 
 
 
 
26a4cfd
f9319c0
 
 
b0efdb8
 
f9319c0
 
 
 
 
b0efdb8
 
 
f9319c0
b0efdb8
 
 
 
 
f9319c0
 
 
 
 
b0efdb8
f9319c0
 
 
 
b0efdb8
f9319c0
b0efdb8
 
f9319c0
 
 
 
 
26a4cfd
f9319c0
 
 
b0efdb8
f9319c0
 
 
b0efdb8
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
b0efdb8
 
 
 
 
f9319c0
 
b0efdb8
26a4cfd
f9319c0
26a4cfd
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
 
b0efdb8
f9319c0
b0efdb8
26a4cfd
f9319c0
 
 
 
 
 
 
 
 
 
 
b0efdb8
 
f9319c0
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
 
 
 
 
 
26a4cfd
b0efdb8
f9319c0
 
 
 
b0efdb8
 
 
26a4cfd
f9319c0
 
 
26a4cfd
f9319c0
 
 
26a4cfd
f9319c0
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
 
 
 
 
 
 
b0efdb8
f9319c0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
f9319c0
ad55c44
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0efdb8
 
 
 
26a4cfd
b0efdb8
e124116
b0efdb8
e124116
b0efdb8
 
f9319c0
b0efdb8
 
26a4cfd
b0efdb8
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
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
import random
import math
import re

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("DiffSketchEdit 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:
        """
        Perform sketch editing using DiffSketchEdit approach
        """
        try:
            # Parse inputs
            if isinstance(inputs, str):
                # Check if it's a JSON string
                try:
                    parsed_inputs = json.loads(inputs)
                    if isinstance(parsed_inputs, dict):
                        inputs = parsed_inputs
                    else:
                        # Simple prompt - treat as generation
                        prompts = [inputs]
                        edit_type = "generate"
                        parameters = {}
                except:
                    # Simple prompt - treat as generation
                    prompts = [inputs]
                    edit_type = "generate"
                    parameters = {}
            
            if isinstance(inputs, dict):
                input_data = inputs.get("inputs", inputs)
                if isinstance(input_data, str):
                    prompts = [input_data]
                    edit_type = "generate"
                elif isinstance(input_data, dict):
                    prompts = input_data.get("prompts", [input_data.get("prompt", "a simple sketch")])
                    edit_type = input_data.get("edit_type", "generate")
                else:
                    prompts = ["a simple sketch"]
                    edit_type = "generate"
                
                parameters = inputs.get("parameters", {})
            
            # Extract parameters with defaults
            width = parameters.get("width", 224)
            height = parameters.get("height", 224)
            seed = parameters.get("seed", None)
            input_svg = parameters.get("input_svg", None)
            
            if seed is not None:
                torch.manual_seed(seed)
                np.random.seed(seed)
                random.seed(seed)
            
            print(f"Processing edit type: '{edit_type}' with prompts: {prompts}")
            
            # Process based on edit type
            if edit_type == "replace" and len(prompts) >= 2:
                svg_content, metadata = self.word_replacement_edit(prompts[0], prompts[1], width, height, input_svg)
            elif edit_type == "refine":
                svg_content, metadata = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
            elif edit_type == "reweight":
                svg_content, metadata = self.attention_reweighting_edit(prompts[0], width, height, input_svg)
            elif edit_type == "generate":
                svg_content, metadata = self.simple_generation(prompts[0], width, height)
            else:
                # Default to refinement
                svg_content, metadata = self.prompt_refinement_edit(prompts[0], width, height, input_svg)
            
            # Convert SVG to PIL Image for HF API compatibility
            pil_image = self.svg_to_pil_image(svg_content, width, height)
            
            # Store metadata
            pil_image.info['svg_content'] = svg_content
            for key, value in metadata.items():
                if isinstance(value, (dict, list)):
                    pil_image.info[key] = json.dumps(value)
                else:
                    pil_image.info[key] = str(value)
            
            return pil_image
            
        except Exception as e:
            print(f"Error in handler: {e}")
            # Return fallback image
            fallback_svg = self.create_fallback_svg(prompts[0] if prompts else "error", width, height)
            fallback_image = self.svg_to_pil_image(fallback_svg, width, height)
            fallback_image.info['error'] = str(e)
            fallback_image.info['edit_type'] = edit_type
            return fallback_image

    def word_replacement_edit(self, source_prompt: str, target_prompt: str, width: int, height: int, input_svg: str = None):
        """Perform word replacement editing"""
        try:
            print(f"Word replacement: '{source_prompt}' -> '{target_prompt}'")
            
            # Analyze word differences
            added_words, removed_words = self.analyze_word_differences(source_prompt, target_prompt)
            print(f"Added words: {added_words}, Removed words: {removed_words}")
            
            # Generate or use base SVG
            if input_svg:
                base_svg = input_svg
            else:
                base_svg = self.generate_base_svg(source_prompt, width, height)
            
            # Apply word replacement transformations
            edited_svg = self.apply_word_replacement(base_svg, source_prompt, target_prompt, added_words, removed_words, width, height)
            
            metadata = {
                "edit_type": "replace",
                "source_prompt": source_prompt,
                "target_prompt": target_prompt,
                "added_words": list(added_words),
                "removed_words": list(removed_words)
            }
            
            return edited_svg, metadata
            
        except Exception as e:
            print(f"Error in word_replacement_edit: {e}")
            fallback_svg = self.create_fallback_svg(source_prompt, width, height)
            metadata = {"edit_type": "replace", "error": str(e)}
            return fallback_svg, metadata

    def prompt_refinement_edit(self, prompt: str, width: int, height: int, input_svg: str = None):
        """Perform prompt refinement editing"""
        try:
            print(f"Prompt refinement for: '{prompt}'")
            
            # Generate or use base SVG
            if input_svg:
                base_svg = input_svg
            else:
                base_svg = self.generate_base_svg(prompt, width, height)
            
            # Apply refinement based on prompt analysis
            refined_svg = self.apply_refinement(base_svg, prompt, width, height)
            
            metadata = {
                "edit_type": "refine",
                "prompt": prompt
            }
            
            return refined_svg, metadata
            
        except Exception as e:
            print(f"Error in prompt_refinement_edit: {e}")
            fallback_svg = self.create_fallback_svg(prompt, width, height)
            metadata = {"edit_type": "refine", "error": str(e)}
            return fallback_svg, metadata

    def attention_reweighting_edit(self, prompt: str, width: int, height: int, input_svg: str = None):
        """Perform attention reweighting editing"""
        try:
            print(f"Attention reweighting for: '{prompt}'")
            
            # Parse attention weights from prompt (e.g., "(cat:1.5)" or "[table:0.5]")
            weighted_prompt, attention_weights = self.parse_attention_weights(prompt)
            print(f"Weighted prompt: '{weighted_prompt}', weights: {attention_weights}")
            
            # Generate or use base SVG
            if input_svg:
                base_svg = input_svg
            else:
                base_svg = self.generate_base_svg(weighted_prompt, width, height)
            
            # Apply attention reweighting
            reweighted_svg = self.apply_attention_reweighting(base_svg, weighted_prompt, attention_weights, width, height)
            
            metadata = {
                "edit_type": "reweight",
                "prompt": prompt,
                "weighted_prompt": weighted_prompt,
                "attention_weights": attention_weights
            }
            
            return reweighted_svg, metadata
            
        except Exception as e:
            print(f"Error in attention_reweighting_edit: {e}")
            fallback_svg = self.create_fallback_svg(prompt, width, height)
            metadata = {"edit_type": "reweight", "error": str(e)}
            return fallback_svg, metadata

    def simple_generation(self, prompt: str, width: int, height: int):
        """Perform simple SVG generation"""
        try:
            print(f"Simple generation for: '{prompt}'")
            
            svg_content = self.generate_base_svg(prompt, width, height)
            
            metadata = {
                "edit_type": "generate",
                "prompt": prompt
            }
            
            return svg_content, metadata
            
        except Exception as e:
            print(f"Error in simple_generation: {e}")
            fallback_svg = self.create_fallback_svg(prompt, width, height)
            metadata = {"edit_type": "generate", "error": str(e)}
            return fallback_svg, metadata

    def generate_base_svg(self, prompt: str, width: int, height: int):
        """Generate base SVG from prompt"""
        dwg = svgwrite.Drawing(size=(width, height))
        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
        
        # Extract semantic features
        features = self.extract_semantic_features(prompt)
        
        # Generate content based on prompt
        if any(word in prompt.lower() for word in ['person', 'people', 'human', 'man', 'woman']):
            self.add_person_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['animal', 'cat', 'dog', 'bird', 'horse']):
            self.add_animal_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['house', 'building', 'architecture']):
            self.add_building_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['tree', 'nature', 'landscape']):
            self.add_nature_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['car', 'vehicle', 'transport']):
            self.add_vehicle_elements(dwg, width, height, features)
        else:
            self.add_abstract_elements(dwg, width, height, features)
        
        return dwg.tostring()

    def analyze_word_differences(self, source: str, target: str):
        """Analyze differences between source and target prompts"""
        source_words = set(source.lower().split())
        target_words = set(target.lower().split())
        
        added_words = target_words - source_words
        removed_words = source_words - target_words
        
        return added_words, removed_words

    def parse_attention_weights(self, prompt: str):
        """Parse attention weights from prompt"""
        # Pattern for (word:weight) - increase attention
        increase_pattern = r'\(([^:]+):([0-9.]+)\)'
        # Pattern for [word:weight] - decrease attention
        decrease_pattern = r'\[([^:]+):([0-9.]+)\]'
        
        attention_weights = {}
        weighted_prompt = prompt
        
        # Find increase weights
        for match in re.finditer(increase_pattern, prompt):
            word = match.group(1).strip()
            weight = float(match.group(2))
            attention_weights[word] = weight
            # Remove the weight notation from prompt
            weighted_prompt = weighted_prompt.replace(match.group(0), word)
        
        # Find decrease weights
        for match in re.finditer(decrease_pattern, prompt):
            word = match.group(1).strip()
            weight = float(match.group(2))
            attention_weights[word] = weight
            # Remove the weight notation from prompt
            weighted_prompt = weighted_prompt.replace(match.group(0), word)
        
        return weighted_prompt.strip(), attention_weights

    def apply_word_replacement(self, base_svg: str, source_prompt: str, target_prompt: str, 
                             added_words: set, removed_words: set, width: int, height: int):
        """Apply word replacement transformations to SVG"""
        # For now, regenerate with target prompt but keep some base structure
        # In a full implementation, this would do more sophisticated editing
        
        # Parse the base SVG to understand its structure
        features = self.extract_semantic_features(target_prompt)
        
        # Create new SVG with target prompt characteristics
        dwg = svgwrite.Drawing(size=(width, height))
        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
        
        # Apply changes based on word differences
        if any(word in added_words for word in ['red', 'blue', 'green', 'yellow']):
            # Color change
            self.add_colored_elements(dwg, width, height, added_words)
        elif any(word in added_words for word in ['big', 'large', 'huge']):
            # Size change
            self.add_large_elements(dwg, width, height, features)
        elif any(word in added_words for word in ['small', 'tiny', 'mini']):
            # Size change
            self.add_small_elements(dwg, width, height, features)
        else:
            # General content change
            self.add_content_based_on_prompt(dwg, target_prompt, width, height)
        
        return dwg.tostring()

    def apply_refinement(self, base_svg: str, prompt: str, width: int, height: int):
        """Apply refinement to existing SVG"""
        # For now, enhance the base SVG with additional details
        features = self.extract_semantic_features(prompt)
        
        dwg = svgwrite.Drawing(size=(width, height))
        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
        
        # Add refined elements based on prompt
        if features.get('detailed', False):
            self.add_detailed_elements(dwg, width, height, features)
        else:
            self.add_content_based_on_prompt(dwg, prompt, width, height)
        
        return dwg.tostring()

    def apply_attention_reweighting(self, base_svg: str, prompt: str, attention_weights: dict, width: int, height: int):
        """Apply attention reweighting to SVG"""
        dwg = svgwrite.Drawing(size=(width, height))
        dwg.add(dwg.rect(insert=(0, 0), size=(width, height), fill='white'))
        
        # Apply different emphasis based on attention weights
        for word, weight in attention_weights.items():
            if weight > 1.0:
                # Emphasize this element
                self.add_emphasized_element(dwg, word, weight, width, height)
            elif weight < 1.0:
                # De-emphasize this element
                self.add_deemphasized_element(dwg, word, weight, width, height)
        
        # Add base content
        self.add_content_based_on_prompt(dwg, prompt, width, height)
        
        return dwg.tostring()

    def add_person_elements(self, dwg, width, height, features):
        """Add person-like elements"""
        center_x, center_y = width // 2, height // 2
        
        # Head
        head_radius = 20
        dwg.add(dwg.circle(center=(center_x, center_y - 40), r=head_radius, fill='#FDBCB4', stroke='black', stroke_width=2))
        
        # Body
        body_height = 60
        body_width = 30
        dwg.add(dwg.rect(
            insert=(center_x - body_width//2, center_y - 10),
            size=(body_width, body_height),
            fill='#4A90E2',
            stroke='black',
            stroke_width=2
        ))
        
        # Arms
        dwg.add(dwg.line(start=(center_x - body_width//2, center_y), end=(center_x - 40, center_y + 20), stroke='black', stroke_width=3))
        dwg.add(dwg.line(start=(center_x + body_width//2, center_y), end=(center_x + 40, center_y + 20), stroke='black', stroke_width=3))
        
        # Legs
        dwg.add(dwg.line(start=(center_x - 10, center_y + body_height - 10), end=(center_x - 20, center_y + body_height + 30), stroke='black', stroke_width=3))
        dwg.add(dwg.line(start=(center_x + 10, center_y + body_height - 10), end=(center_x + 20, center_y + body_height + 30), stroke='black', stroke_width=3))

    def add_animal_elements(self, dwg, width, height, features):
        """Add animal-like elements"""
        center_x, center_y = width // 2, height // 2
        
        # Body (oval)
        dwg.add(dwg.ellipse(center=(center_x, center_y), r=(40, 25), fill='#8B4513', stroke='black', stroke_width=2))
        
        # Head
        dwg.add(dwg.circle(center=(center_x - 30, center_y - 10), r=20, fill='#A0522D', stroke='black', stroke_width=2))
        
        # Legs
        for i, x_offset in enumerate([-20, -10, 10, 20]):
            dwg.add(dwg.line(
                start=(center_x + x_offset, center_y + 25),
                end=(center_x + x_offset, center_y + 45),
                stroke='black',
                stroke_width=3
            ))
        
        # Tail
        dwg.add(dwg.path(
            d=f"M {center_x + 40},{center_y} Q {center_x + 60},{center_y - 20} {center_x + 50},{center_y - 35}",
            stroke='black',
            stroke_width=3,
            fill='none'
        ))

    def add_building_elements(self, dwg, width, height, features):
        """Add building-like elements"""
        # Main building
        building_width = width * 0.6
        building_height = height * 0.7
        x = (width - building_width) // 2
        y = height - building_height - 10
        
        dwg.add(dwg.rect(
            insert=(x, y),
            size=(building_width, building_height),
            fill='#CD853F',
            stroke='black',
            stroke_width=2
        ))
        
        # Roof
        roof_points = [(x, y), (x + building_width//2, y - 30), (x + building_width, y)]
        dwg.add(dwg.polygon(points=roof_points, fill='#8B0000', stroke='black', stroke_width=2))
        
        # Windows
        window_size = 15
        for i in range(3):
            for j in range(4):
                wx = x + 15 + i * 30
                wy = y + 15 + j * 25
                if wy < y + building_height - 20:
                    dwg.add(dwg.rect(
                        insert=(wx, wy),
                        size=(window_size, window_size),
                        fill='#87CEEB',
                        stroke='black',
                        stroke_width=1
                    ))
        
        # Door
        door_width = 20
        door_height = 40
        door_x = x + building_width//2 - door_width//2
        door_y = y + building_height - door_height
        dwg.add(dwg.rect(
            insert=(door_x, door_y),
            size=(door_width, door_height),
            fill='#8B4513',
            stroke='black',
            stroke_width=2
        ))

    def add_nature_elements(self, dwg, width, height, features):
        """Add nature-like elements"""
        # Tree
        center_x, center_y = width // 2, height // 2
        
        # Trunk
        trunk_width = 15
        trunk_height = height // 3
        trunk_x = center_x - trunk_width // 2
        trunk_y = height - trunk_height - 10
        
        dwg.add(dwg.rect(
            insert=(trunk_x, trunk_y),
            size=(trunk_width, trunk_height),
            fill='#8B4513',
            stroke='black',
            stroke_width=1
        ))
        
        # Crown (multiple circles for foliage)
        crown_radius = 30
        for i, (dx, dy) in enumerate([(-15, -20), (15, -20), (0, -35), (-10, -50), (10, -50)]):
            dwg.add(dwg.circle(
                center=(center_x + dx, center_y + dy),
                r=crown_radius - i * 3,
                fill='#228B22',
                stroke='#006400',
                stroke_width=1,
                opacity=0.8
            ))

    def add_vehicle_elements(self, dwg, width, height, features):
        """Add vehicle-like elements"""
        center_x, center_y = width // 2, height // 2
        
        # Car body
        car_width = width * 0.6
        car_height = height * 0.3
        car_x = (width - car_width) // 2
        car_y = center_y + 10
        
        dwg.add(dwg.rect(
            insert=(car_x, car_y),
            size=(car_width, car_height),
            fill='#FF4500',
            stroke='black',
            stroke_width=2,
            rx=5
        ))
        
        # Windshield
        windshield_width = car_width * 0.6
        windshield_height = car_height * 0.4
        windshield_x = car_x + (car_width - windshield_width) // 2
        windshield_y = car_y - windshield_height + 5
        
        dwg.add(dwg.rect(
            insert=(windshield_x, windshield_y),
            size=(windshield_width, windshield_height),
            fill='#87CEEB',
            stroke='black',
            stroke_width=1
        ))
        
        # Wheels
        wheel_radius = 12
        wheel_y = car_y + car_height - 5
        dwg.add(dwg.circle(center=(car_x + 25, wheel_y), r=wheel_radius, fill='black'))
        dwg.add(dwg.circle(center=(car_x + car_width - 25, wheel_y), r=wheel_radius, fill='black'))

    def add_abstract_elements(self, dwg, width, height, features):
        """Add abstract elements"""
        colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#96CEB4', '#FFEAA7']
        
        for i in range(5):
            shape_type = random.choice(['circle', 'rect', 'path'])
            color = random.choice(colors)
            
            if shape_type == 'circle':
                radius = random.randint(10, 30)
                x = random.randint(radius, width - radius)
                y = random.randint(radius, height - radius)
                dwg.add(dwg.circle(center=(x, y), r=radius, fill=color, opacity=0.7))
            elif shape_type == 'rect':
                w = random.randint(20, 60)
                h = random.randint(20, 60)
                x = random.randint(0, width - w)
                y = random.randint(0, height - h)
                dwg.add(dwg.rect(insert=(x, y), size=(w, h), fill=color, opacity=0.7))
            else:
                # Random path
                start_x = random.randint(0, width)
                start_y = random.randint(0, height)
                end_x = random.randint(0, width)
                end_y = random.randint(0, height)
                dwg.add(dwg.line(start=(start_x, start_y), end=(end_x, end_y), stroke=color, stroke_width=3))

    def add_colored_elements(self, dwg, width, height, color_words):
        """Add elements with specific colors"""
        color_map = {
            'red': '#FF0000',
            'blue': '#0000FF',
            'green': '#00FF00',
            'yellow': '#FFFF00',
            'purple': '#800080',
            'orange': '#FFA500'
        }
        
        center_x, center_y = width // 2, height // 2
        
        for word in color_words:
            if word in color_map:
                color = color_map[word]
                # Add a colored shape
                dwg.add(dwg.circle(
                    center=(center_x + random.randint(-50, 50), center_y + random.randint(-50, 50)),
                    r=random.randint(15, 35),
                    fill=color,
                    opacity=0.8
                ))

    def add_large_elements(self, dwg, width, height, features):
        """Add large-sized elements"""
        center_x, center_y = width // 2, height // 2
        
        # Large central element
        dwg.add(dwg.circle(
            center=(center_x, center_y),
            r=min(width, height) // 3,
            fill='#4A90E2',
            stroke='black',
            stroke_width=3
        ))

    def add_small_elements(self, dwg, width, height, features):
        """Add small-sized elements"""
        # Multiple small elements
        for i in range(8):
            x = random.randint(10, width - 10)
            y = random.randint(10, height - 10)
            dwg.add(dwg.circle(
                center=(x, y),
                r=random.randint(3, 8),
                fill='#E74C3C',
                opacity=0.7
            ))

    def add_detailed_elements(self, dwg, width, height, features):
        """Add detailed elements for refinement"""
        # Add more complex shapes and details
        self.add_abstract_elements(dwg, width, height, features)
        
        # Add decorative elements
        center_x, center_y = width // 2, height // 2
        for i in range(4):
            angle = i * math.pi / 2
            x = center_x + 40 * math.cos(angle)
            y = center_y + 40 * math.sin(angle)
            dwg.add(dwg.circle(center=(x, y), r=8, fill='#9B59B6', opacity=0.6))

    def add_emphasized_element(self, dwg, word: str, weight: float, width: int, height: int):
        """Add emphasized element based on attention weight"""
        center_x, center_y = width // 2, height // 2
        
        # Scale size based on weight
        base_size = 20
        size = int(base_size * weight)
        
        dwg.add(dwg.circle(
            center=(center_x + random.randint(-30, 30), center_y + random.randint(-30, 30)),
            r=size,
            fill='#FF6B6B',
            opacity=min(1.0, weight / 2),
            stroke='black',
            stroke_width=2
        ))

    def add_deemphasized_element(self, dwg, word: str, weight: float, width: int, height: int):
        """Add de-emphasized element based on attention weight"""
        center_x, center_y = width // 2, height // 2
        
        # Scale size based on weight
        base_size = 15
        size = int(base_size * weight)
        
        dwg.add(dwg.circle(
            center=(center_x + random.randint(-40, 40), center_y + random.randint(-40, 40)),
            r=max(3, size),
            fill='#CCCCCC',
            opacity=weight,
            stroke='gray',
            stroke_width=1
        ))

    def add_content_based_on_prompt(self, dwg, prompt: str, width: int, height: int):
        """Add content based on prompt analysis"""
        features = self.extract_semantic_features(prompt)
        
        if any(word in prompt.lower() for word in ['person', 'people', 'human']):
            self.add_person_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['animal', 'cat', 'dog']):
            self.add_animal_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['house', 'building']):
            self.add_building_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['tree', 'nature']):
            self.add_nature_elements(dwg, width, height, features)
        elif any(word in prompt.lower() for word in ['car', 'vehicle']):
            self.add_vehicle_elements(dwg, width, height, features)
        else:
            self.add_abstract_elements(dwg, width, height, features)

    def extract_semantic_features(self, prompt: str):
        """Extract semantic features from prompt"""
        features = {
            'detailed': False,
            'simple': False,
            'colorful': False,
            'large': False,
            'small': False
        }
        
        prompt_lower = prompt.lower()
        
        if any(word in prompt_lower for word in ['detailed', 'complex', 'intricate']):
            features['detailed'] = True
        if any(word in prompt_lower for word in ['simple', 'minimal', 'basic']):
            features['simple'] = True
        if any(word in prompt_lower for word in ['colorful', 'bright', 'vibrant']):
            features['colorful'] = True
        if any(word in prompt_lower for word in ['large', 'big', 'huge']):
            features['large'] = True
        if any(word in prompt_lower for word in ['small', 'tiny', 'mini']):
            features['small'] = True
        
        return features

    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
        prompt_str = str(prompt)[:30] if prompt else "error"
        dwg.add(dwg.text(
            f"DiffSketchEdit\n{prompt_str}...",
            insert=(width/2, height/2),
            text_anchor="middle",
            font_size="12px",
            fill="black"
        ))
        
        return dwg.tostring()