primerz commited on
Commit
d432eb2
·
verified ·
1 Parent(s): b867149

Update generator.py

Browse files
Files changed (1) hide show
  1. generator.py +201 -512
generator.py CHANGED
@@ -1,5 +1,6 @@
1
  """
2
  Generation logic for Pixagram AI Pixel Art Generator
 
3
  """
4
  import torch
5
  import numpy as np
@@ -14,17 +15,18 @@ from config import (
14
  )
15
  from utils import (
16
  sanitize_text, enhanced_color_match, color_match, create_face_mask,
17
- draw_kps, get_demographic_description, calculate_optimal_size, enhance_face_crop
 
18
  )
19
  from models import (
20
- load_face_analysis, load_depth_detector, load_controlnets, load_image_encoder,
21
- load_sdxl_pipeline, load_lora, setup_ip_adapter, setup_compel,
22
  setup_scheduler, optimize_pipeline, load_caption_model, set_clip_skip
23
  )
24
 
25
 
26
  class RetroArtConverter:
27
- """Main class for retro art generation"""
28
 
29
  def __init__(self):
30
  self.device = device
@@ -33,8 +35,7 @@ class RetroArtConverter:
33
  'custom_checkpoint': False,
34
  'lora': False,
35
  'instantid': False,
36
- 'zoe_depth': False,
37
- 'ip_adapter': False
38
  }
39
 
40
  # Initialize face analysis
@@ -44,27 +45,14 @@ class RetroArtConverter:
44
  self.zoe_depth, zoe_success = load_depth_detector()
45
  self.models_loaded['zoe_depth'] = zoe_success
46
 
47
- # Load ControlNets
48
- controlnet_depth, self.controlnet_instantid, instantid_success = load_controlnets()
49
- self.controlnet_depth = controlnet_depth
50
- self.instantid_enabled = instantid_success
51
- self.models_loaded['instantid'] = instantid_success
52
 
53
- # Load image encoder
54
- if self.instantid_enabled:
55
- self.image_encoder = load_image_encoder()
56
- else:
57
- self.image_encoder = None
58
-
59
- # Determine which controlnets to use
60
- if self.instantid_enabled and self.controlnet_instantid is not None:
61
- controlnets = [self.controlnet_instantid, controlnet_depth]
62
- print(f"Initializing with multiple ControlNets: InstantID + Depth")
63
- else:
64
- controlnets = controlnet_depth
65
- print(f"Initializing with single ControlNet: Depth only")
66
 
67
- # Load SDXL pipeline
68
  self.pipe, checkpoint_success = load_sdxl_pipeline(controlnets)
69
  self.models_loaded['custom_checkpoint'] = checkpoint_success
70
 
@@ -72,15 +60,6 @@ class RetroArtConverter:
72
  lora_success = load_lora(self.pipe)
73
  self.models_loaded['lora'] = lora_success
74
 
75
- # Setup IP-Adapter
76
- if self.instantid_enabled and self.image_encoder is not None:
77
- self.image_proj_model, ip_adapter_success = setup_ip_adapter(self.pipe, self.image_encoder)
78
- self.models_loaded['ip_adapter'] = ip_adapter_success
79
- else:
80
- print("[INFO] Face preservation: InstantID ControlNet keypoints only")
81
- self.models_loaded['ip_adapter'] = False
82
- self.image_proj_model = None
83
-
84
  # Setup Compel
85
  self.compel, self.use_compel = setup_compel(self.pipe)
86
 
@@ -99,21 +78,14 @@ class RetroArtConverter:
99
  print(" [OK] Using GIT for detailed captions")
100
  elif self.caption_model_type == "blip":
101
  print(" [OK] Using BLIP for standard captions")
102
- else:
103
- print(" [OK] Caption model loaded")
104
-
105
 
106
  # Set CLIP skip
107
  set_clip_skip(self.pipe)
108
 
109
- # Track controlnet configuration
110
- self.using_multiple_controlnets = isinstance(controlnets, list)
111
- print(f"Pipeline initialized with {'multiple' if self.using_multiple_controlnets else 'single'} ControlNet(s)")
112
-
113
  # Print model status
114
  self._print_status()
115
 
116
- print(" [OK] Model initialization complete!")
117
 
118
  def _print_status(self):
119
  """Print model loading status"""
@@ -121,309 +93,82 @@ class RetroArtConverter:
121
  for model, loaded in self.models_loaded.items():
122
  status = "[OK] LOADED" if loaded else "[FALLBACK/DISABLED]"
123
  print(f"{model}: {status}")
 
124
  print("===================\n")
125
-
126
- print("=== UPGRADE VERIFICATION ===")
127
- try:
128
- from resampler_enhanced import EnhancedResampler
129
- from ip_attention_processor_enhanced import EnhancedIPAttnProcessor2_0
130
-
131
- resampler_check = isinstance(self.image_proj_model, EnhancedResampler) if hasattr(self, 'image_proj_model') and self.image_proj_model is not None else False
132
- custom_attn_check = any(isinstance(p, EnhancedIPAttnProcessor2_0) for p in self.pipe.unet.attn_processors.values()) if hasattr(self, 'pipe') else False
133
-
134
- print(f"Enhanced Perceiver Resampler: {'[OK] ACTIVE' if resampler_check else '[INFO] Not active'}")
135
- print(f"Enhanced IP-Adapter Attention: {'[OK] ACTIVE' if custom_attn_check else '[INFO] Not active'}")
136
-
137
- if resampler_check and custom_attn_check:
138
- print("[SUCCESS] Face preservation upgrade fully active")
139
- print(" Expected improvement: +10-15% face similarity")
140
- elif resampler_check or custom_attn_check:
141
- print("[PARTIAL] Some upgrades active")
142
- else:
143
- print("[INFO] Using standard components")
144
- except Exception as e:
145
- print(f"[INFO] Verification skipped: {e}")
146
- print("============================\n")
147
 
148
  def get_depth_map(self, image):
149
- """Generate depth map using Zoe Depth"""
150
- if self.zoe_depth is not None:
151
- try:
152
- if image.mode != 'RGB':
153
- image = image.convert('RGB')
154
-
155
- orig_width, orig_height = image.size
156
- # **FIX 1 START: Ensure all size variables are standard Python int**
157
- orig_width = int(orig_width)
158
- orig_height = int(orig_height)
159
-
160
- # FIXED: Use multiples of 64 (not 32)
161
- target_width = int((orig_width // 64) * 64)
162
- target_height = int((orig_height // 64) * 64)
163
-
164
- target_width = int(max(64, target_width))
165
- target_height = int(max(64, target_height))
166
-
167
- # Create an explicit tuple of standard ints
168
- size_for_depth = (int(target_width), int(target_height))
169
-
170
- # Always resize using the explicit int tuple to avoid numpy.int64 issues
171
- # This replaces the conditional resize
172
- image_for_depth = image.resize(size_for_depth, Image.LANCZOS)
173
-
174
- if target_width != orig_width or target_height != orig_height:
175
- print(f"[DEPTH] Resized for ZoeDetector: {orig_width}x{orig_height} -> {target_width}x{target_height}")
176
-
177
- # FIXED: Add torch.no_grad() wrapper
178
- with torch.no_grad():
179
- depth_image = self.zoe_depth(image_for_depth) # Use the correctly-typed resized image
180
-
181
- depth_width, depth_height = depth_image.size
182
- if depth_width != orig_width or depth_height != orig_height:
183
- # Resize back to the original size that get_depth_map received
184
- depth_image = depth_image.resize((int(orig_width), int(orig_height)), Image.LANCZOS)
185
- # **FIX 1 END**
186
-
187
- print(f"[DEPTH] Zoe depth map generated: {orig_width}x{orig_height}")
188
- return depth_image
189
-
190
- except Exception as e:
191
- print(f"[DEPTH] ZoeDetector failed ({e}), falling back to grayscale depth")
192
- gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
193
- depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
194
- return Image.fromarray(depth_colored)
195
- else:
196
- gray = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2GRAY)
197
- depth_colored = cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
198
- return Image.fromarray(depth_colored)
199
-
200
-
201
- def add_trigger_word(self, prompt):
202
- """Add trigger word to prompt if not present"""
203
- if TRIGGER_WORD.lower() not in prompt.lower():
204
- # **FIX 3 START: Handle empty or blank prompt**
205
- if not prompt or not prompt.strip():
206
- return TRIGGER_WORD
207
- # **FIX 3 END**
208
- return f"{TRIGGER_WORD}, {prompt}"
209
- return prompt
210
-
211
- def extract_multi_scale_face(self, face_crop, face):
212
- """
213
- Extract face features at multiple scales for better detail.
214
- +1-2% improvement in face preservation.
215
- """
216
- try:
217
- multi_scale_embeds = []
218
-
219
- for scale in MULTI_SCALE_FACTORS:
220
- # Resize
221
- w, h = face_crop.size
222
- scaled_size = (int(w * scale), int(h * scale))
223
- scaled_crop = face_crop.resize(scaled_size, Image.LANCZOS)
224
 
225
- # Pad/crop back to original
226
- scaled_crop = scaled_crop.resize((w, h), Image.LANCZOS)
227
 
228
- # Extract features
229
- scaled_array = cv2.cvtColor(np.array(scaled_crop), cv2.COLOR_RGB2BGR)
230
- scaled_faces = self.face_app.get(scaled_array)
231
 
232
- if len(scaled_faces) > 0:
233
- multi_scale_embeds.append(scaled_faces[0].normed_embedding)
234
-
235
- # Average embeddings
236
- if len(multi_scale_embeds) > 0:
237
- averaged = np.mean(multi_scale_embeds, axis=0)
238
- # Renormalize
239
- averaged = averaged / np.linalg.norm(averaged)
240
- print(f"[MULTI-SCALE] Combined {len(multi_scale_embeds)} scales")
241
- return averaged
242
-
243
- return face.normed_embedding
244
-
245
- except Exception as e:
246
- print(f"[MULTI-SCALE] Failed: {e}, using single scale")
247
- return face.normed_embedding
248
-
249
- def detect_face_quality(self, face):
250
- """
251
- Detect face quality and adaptively adjust parameters.
252
- +2-3% consistency improvement.
253
- """
254
- try:
255
- bbox = face.bbox
256
- face_size = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
257
- det_score = float(face.det_score) if hasattr(face, 'det_score') else 1.0
258
-
259
- # Small face -> boost identity preservation
260
- if face_size < ADAPTIVE_THRESHOLDS['small_face_size']:
261
- return ADAPTIVE_PARAMS['small_face'].copy()
262
-
263
- # Low confidence -> boost preservation
264
- elif det_score < ADAPTIVE_THRESHOLDS['low_confidence']:
265
- return ADAPTIVE_PARAMS['low_confidence'].copy()
266
-
267
- # Check for profile/side view (if pose available)
268
- elif hasattr(face, 'pose') and len(face.pose) > 1:
269
- try:
270
- yaw = float(face.pose[1])
271
- if abs(yaw) > ADAPTIVE_THRESHOLDS['profile_angle']:
272
- return ADAPTIVE_PARAMS['profile_view'].copy()
273
- except (ValueError, TypeError, IndexError):
274
- pass
275
-
276
- # Good quality face - use provided parameters
277
- return None
278
-
279
- except Exception as e:
280
- print(f"[ADAPTIVE] Quality detection failed: {e}")
281
- return None
282
-
283
- def validate_and_adjust_parameters(self, strength, guidance_scale, lora_scale,
284
- identity_preservation, identity_control_scale,
285
- depth_control_scale, consistency_mode=True):
286
- """
287
- Enhanced parameter validation with stricter rules for consistency.
288
- """
289
- if consistency_mode:
290
- print("[CONSISTENCY] Applying strict parameter validation...")
291
- adjustments = []
292
-
293
- # Rule 1: Strong inverse relationship between identity and LORA
294
- if identity_preservation > 1.2:
295
- original_lora = lora_scale
296
- lora_scale = min(lora_scale, 1.0)
297
- if abs(lora_scale - original_lora) > 0.01:
298
- adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (high identity)")
299
-
300
- # Rule 2: Strength-based profile activation
301
- if strength < 0.5:
302
- # Maximum preservation mode
303
- if identity_preservation < 1.3:
304
- original_identity = identity_preservation
305
- identity_preservation = 1.3
306
- adjustments.append(f"Identity: {original_identity:.2f}->{identity_preservation:.2f} (max preservation)")
307
- if lora_scale > 0.9:
308
- original_lora = lora_scale
309
- lora_scale = 0.9
310
- adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (max preservation)")
311
- if guidance_scale > 1.3:
312
- original_cfg = guidance_scale
313
- guidance_scale = 1.3
314
- adjustments.append(f"CFG: {original_cfg:.2f}->{guidance_scale:.2f} (max preservation)")
315
-
316
- elif strength > 0.7:
317
- # Artistic transformation mode
318
- if identity_preservation > 1.0:
319
- original_identity = identity_preservation
320
- identity_preservation = 1.0
321
- adjustments.append(f"Identity: {original_identity:.2f}->{identity_preservation:.2f} (artistic mode)")
322
- if lora_scale < 1.2:
323
- original_lora = lora_scale
324
- lora_scale = 1.2
325
- adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (artistic mode)")
326
-
327
- # Rule 3: CFG-LORA relationship
328
- if guidance_scale > 1.4 and lora_scale > 1.2:
329
- original_lora = lora_scale
330
- lora_scale = 1.1
331
- adjustments.append(f"LORA: {original_lora:.2f}->{lora_scale:.2f} (high CFG detected)")
332
-
333
- # Rule 4: LCM sweet spot enforcement
334
- original_cfg = guidance_scale
335
- guidance_scale = max(1.0, min(guidance_scale, 1.5))
336
- if abs(guidance_scale - original_cfg) > 0.01:
337
- adjustments.append(f"CFG: {original_cfg:.2f}->{guidance_scale:.2f} (LCM optimal)")
338
-
339
- # Rule 5: ControlNet balance
340
- total_control = identity_control_scale + depth_control_scale
341
- if total_control > 1.7:
342
- scale_factor = 1.7 / total_control
343
- original_id_ctrl = identity_control_scale
344
- original_depth_ctrl = depth_control_scale
345
- identity_control_scale *= scale_factor
346
- depth_control_scale *= scale_factor
347
- adjustments.append(f"ControlNets balanced: ID {original_id_ctrl:.2f}->{identity_control_scale:.2f}, Depth {original_depth_ctrl:.2f}->{depth_control_scale:.2f}")
348
-
349
- # Report adjustments
350
- if adjustments:
351
- print(" [OK] Applied adjustments:")
352
- for adj in adjustments:
353
- print(f" - {adj}")
354
- else:
355
- print(" [OK] Parameters already optimal")
356
-
357
- return strength, guidance_scale, lora_scale, identity_preservation, identity_control_scale, depth_control_scale
358
-
359
- def generate_caption(self, image, max_length=None, num_beams=None):
360
- """Generate a descriptive caption for the image (supports BLIP-2, GIT, BLIP)."""
361
- if not self.caption_enabled or self.caption_model is None:
362
- return None
363
-
364
- # Set defaults based on model type
365
- if max_length is None:
366
- if self.caption_model_type == "blip2":
367
- max_length = 50 # BLIP-2 can handle longer captions
368
- elif self.caption_model_type == "git":
369
- max_length = 40 # GIT also produces good long captions
370
- else:
371
- max_length = CAPTION_CONFIG['max_length'] # BLIP base (20)
372
-
373
- if num_beams is None:
374
- num_beams = CAPTION_CONFIG['num_beams']
375
-
376
- try:
377
- if self.caption_model_type == "blip2":
378
- # BLIP-2 specific processing
379
- inputs = self.caption_processor(image, return_tensors="pt").to(self.device, self.dtype)
380
 
381
- with torch.no_grad():
382
- output = self.caption_model.generate(
383
- **inputs,
384
- max_length=max_length,
385
- num_beams=num_beams,
386
- min_length=10, # Encourage longer captions
387
- length_penalty=1.0,
388
- repetition_penalty=1.5,
389
- early_stopping=True
390
- )
391
 
392
- caption = self.caption_processor.decode(output[0], skip_special_tokens=True)
 
393
 
394
- elif self.caption_model_type == "git":
395
- # GIT specific processing
396
- inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device, self.dtype)
397
 
398
- with torch.no_grad():
399
- output = self.caption_model.generate(
400
- pixel_values=inputs.pixel_values,
401
- max_length=max_length,
402
- num_beams=num_beams,
403
- min_length=10,
404
- length_penalty=1.0,
405
- repetition_penalty=1.5,
406
- early_stopping=True
407
- )
408
 
409
- caption = self.caption_processor.batch_decode(output, skip_special_tokens=True)[0]
 
 
410
 
411
- else:
412
- # BLIP base processing
413
- inputs = self.caption_processor(image, return_tensors="pt").to(self.device, self.dtype)
414
 
415
- with torch.no_grad():
416
- output = self.caption_model.generate(
417
- **inputs,
418
- max_length=max_length,
419
- num_beams=num_beams,
420
- early_stopping=True
421
- )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
422
 
423
- caption = self.caption_processor.decode(output[0], skip_special_tokens=True)
 
 
 
 
 
 
 
 
 
424
 
425
- return caption.strip()
426
-
427
  except Exception as e:
428
  print(f"Caption generation failed: {e}")
429
  return None
@@ -431,134 +176,111 @@ class RetroArtConverter:
431
  def generate_retro_art(
432
  self,
433
  input_image,
434
- prompt="retro game character, vibrant colors, detailed",
435
- negative_prompt="blurry, low quality, ugly, distorted",
436
  num_inference_steps=12,
437
- guidance_scale=1.0,
438
- depth_control_scale=0.8,
439
  identity_control_scale=0.85,
440
  lora_scale=1.0,
441
- identity_preservation=0.8,
442
- strength=0.75,
443
  enable_color_matching=False,
444
  consistency_mode=True,
445
  seed=-1
446
  ):
447
- """Generate retro art with img2img pipeline and enhanced InstantID"""
448
-
449
- # Sanitize text inputs
450
- prompt = sanitize_text(prompt)
451
- negative_prompt = sanitize_text(negative_prompt)
452
 
453
- # **FIX 3 START: Ensure blank negative prompts are empty strings for Compel**
454
- if not negative_prompt or not negative_prompt.strip():
455
- negative_prompt = ""
456
- # **FIX 3 END**
457
 
458
- # Apply parameter validation
459
  if consistency_mode:
460
- print("\n[CONSISTENCY] Validating and adjusting parameters...")
461
- strength, guidance_scale, lora_scale, identity_preservation, identity_control_scale, depth_control_scale = \
462
- self.validate_and_adjust_parameters(
463
- strength, guidance_scale, lora_scale, identity_preservation,
464
- identity_control_scale, depth_control_scale, consistency_mode
465
- )
466
-
467
- # Add trigger word (handles blank prompt fix)
468
- prompt = self.add_trigger_word(prompt)
469
-
470
- # Calculate optimal size with flexible aspect ratio support
471
- original_width, original_height = input_image.size
472
- target_width, target_height = calculate_optimal_size(original_width, original_height)
 
 
473
 
474
- print(f"Resizing from {original_width}x{original_height} to {target_width}x{target_height}")
475
- print(f"Prompt: {prompt}")
476
- print(f"Img2Img Strength: {strength}")
477
 
478
- # Resize with high quality
479
- resized_image = input_image.resize((int(target_width), int(target_height)), Image.LANCZOS)
480
 
481
  # Generate depth map
482
- print("Generating Zoe depth map...")
483
- depth_image = self.get_depth_map(resized_image)
484
- if depth_image.size != (target_width, target_height):
485
- depth_image = depth_image.resize((int(target_width), int(target_height)), Image.LANCZOS)
 
486
 
487
- # Handle face detection
488
- using_multiple_controlnets = self.using_multiple_controlnets
 
489
  face_kps_image = None
490
  face_embeddings = None
491
- face_crop_enhanced = None
492
- has_detected_faces = False
493
  face_bbox_original = None
494
 
495
- if using_multiple_controlnets and self.face_app is not None:
496
- print("Detecting faces and extracting keypoints...")
497
- img_array = cv2.cvtColor(np.array(resized_image), cv2.COLOR_RGB2BGR)
498
- faces = self.face_app.get(img_array)
499
-
500
- if len(faces) > 0:
501
- has_detected_faces = True
502
- print(f"Detected {len(faces)} face(s)")
503
-
504
- # Get largest face
505
- face = sorted(faces, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))[-1]
506
-
507
- # ADAPTIVE PARAMETERS
508
- adaptive_params = self.detect_face_quality(face)
509
- if adaptive_params is not None:
510
- print(f"[ADAPTIVE] {adaptive_params['reason']}")
511
- identity_preservation = adaptive_params['identity_preservation']
512
- identity_control_scale = adaptive_params['identity_control_scale']
513
- guidance_scale = adaptive_params['guidance_scale']
514
- lora_scale = adaptive_params['lora_scale']
515
-
516
- # Extract face embeddings
517
- face_embeddings_base = face.normed_embedding
518
-
519
- # Extract face crop
520
- bbox = face.bbox.astype(int)
521
- x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
522
- face_bbox_original = [x1, y1, x2, y2]
523
-
524
- # Add padding
525
- face_width = x2 - x1
526
- face_height = y2 - y1
527
- padding_x = int(face_width * 0.3)
528
- padding_y = int(face_height * 0.3)
529
- x1 = max(0, x1 - padding_x)
530
- y1 = max(0, y1 - padding_y)
531
- x2 = min(resized_image.width, x2 + padding_x)
532
- y2 = min(resized_image.height, y2 + padding_y)
533
-
534
- # Crop face region
535
- face_crop = resized_image.crop((x1, y1, x2, y2))
536
-
537
- # MULTI-SCALE PROCESSING
538
- face_embeddings = self.extract_multi_scale_face(face_crop, face)
539
-
540
- # Enhance face crop
541
- face_crop_enhanced = enhance_face_crop(face_crop)
542
-
543
- # Draw keypoints
544
- face_kps = face.kps
545
- face_kps_image = draw_kps(resized_image, face_kps)
546
-
547
- # ENHANCED: Extract comprehensive facial attributes
548
- from utils import get_facial_attributes, build_enhanced_prompt
549
- facial_attrs = get_facial_attributes(face)
550
-
551
- # Update prompt with detected attributes
552
- prompt = build_enhanced_prompt(prompt, facial_attrs, TRIGGER_WORD)
553
-
554
- # Legacy output for compatibility
555
- age = facial_attrs['age']
556
- gender_code = facial_attrs['gender']
557
- det_score = facial_attrs['quality']
558
 
559
- gender_str = 'M' if gender_code == 1 else ('F' if gender_code == 0 else 'N/A')
560
- print(f"Face info: bbox={face.bbox}, age={age if age else 'N/A'}, gender={gender_str}")
561
- print(f"Face crop size: {face_crop.size}, enhanced: {face_crop_enhanced.size if face_crop_enhanced else 'N/A'}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
562
 
563
  # Set LORA scale
564
  if hasattr(self.pipe, 'set_adapters') and self.models_loaded['lora']:
@@ -613,91 +335,54 @@ class RetroArtConverter:
613
  if hasattr(self.pipe, 'text_encoder'):
614
  pipe_kwargs["clip_skip"] = 2
615
 
616
- # Configure ControlNet inputs
617
- if using_multiple_controlnets and has_detected_faces and face_kps_image is not None:
618
- print("Using InstantID (keypoints) + Depth ControlNets")
619
- control_images = [face_kps_image, depth_image]
620
- conditioning_scales = [identity_control_scale, depth_control_scale]
621
 
622
- pipe_kwargs["control_image"] = control_images
623
- pipe_kwargs["controlnet_conditioning_scale"] = conditioning_scales
624
-
625
- # Add face embeddings for IP-Adapter if available
626
- if face_embeddings is not None and self.models_loaded.get('ip_adapter', False) and face_crop_enhanced is not None:
627
- print(f"Processing InstantID face embeddings with Resampler...")
628
-
629
- with torch.no_grad():
630
- # Convert InsightFace embeddings to tensor
631
- face_emb_tensor = torch.from_numpy(face_embeddings).to(
632
- device=self.device,
633
- dtype=self.dtype
634
- )
635
-
636
- # Reshape for Resampler: [1, 1, 512]
637
- face_emb_tensor = face_emb_tensor.reshape(1, -1, 512)
638
-
639
- # Pass through Resampler: [1, 1, 512] → [1, 16, 2048]
640
- face_proj_embeds = self.image_proj_model(face_emb_tensor)
641
-
642
- # Scale with identity preservation
643
- boosted_scale = identity_preservation * IDENTITY_BOOST_MULTIPLIER
644
- face_proj_embeds = face_proj_embeds * boosted_scale
645
-
646
- print(f" - Face embedding: {face_emb_tensor.shape}")
647
- print(f" - Resampler output: {face_proj_embeds.shape}")
648
- print(f" - Scale: {boosted_scale:.2f}")
649
-
650
- # CRITICAL: Concatenate with text embeddings (not separate kwargs!)
651
- if 'prompt_embeds' in pipe_kwargs:
652
- # Compel encoded prompts
653
- original_embeds = pipe_kwargs['prompt_embeds']
654
-
655
- # Handle CFG (classifier-free guidance)
656
- if original_embeds.shape[0] > 1: # Has negative + positive
657
- # Duplicate for negative + positive
658
- face_proj_embeds = torch.cat([
659
- torch.zeros_like(face_proj_embeds), # Negative
660
- face_proj_embeds # Positive
661
- ], dim=0)
662
-
663
- # Concatenate: [batch, text_tokens, 2048] + [batch, 16, 2048]
664
- combined_embeds = torch.cat([original_embeds, face_proj_embeds], dim=1)
665
- pipe_kwargs['prompt_embeds'] = combined_embeds
666
-
667
- print(f" - Text embeds: {original_embeds.shape}")
668
- print(f" - Combined embeds: {combined_embeds.shape}")
669
- print(f" [OK] Face embeddings concatenated successfully!")
670
-
671
- else:
672
- print(f" [WARNING] Can't concatenate - no prompt_embeds (use Compel)")
673
 
674
- elif has_detected_faces and self.models_loaded.get('ip_adapter', False):
675
- # Face detected but embeddings unavailable
676
- print(" Face detected but embeddings unavailable, using keypoints only")
677
- # No need for dummy embeddings with concatenation approach
678
-
679
- elif using_multiple_controlnets and not has_detected_faces:
680
- print("Multiple ControlNets available but no faces detected, using depth only")
681
- control_images = [depth_image, depth_image]
682
- conditioning_scales = [0.0, depth_control_scale]
683
 
684
- pipe_kwargs["control_image"] = control_images
685
- pipe_kwargs["controlnet_conditioning_scale"] = conditioning_scales
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
686
 
687
  else:
688
- print("Using Depth ControlNet only")
689
- pipe_kwargs["control_image"] = depth_image
690
- pipe_kwargs["controlnet_conditioning_scale"] = depth_control_scale
691
-
 
692
 
693
  # Generate
694
  print(f"Generating with LCM: Steps={num_inference_steps}, CFG={guidance_scale}, Strength={strength}")
695
- print(f"Controlnet scales - Identity: {identity_control_scale}, Depth: {depth_control_scale}")
696
  result = self.pipe(**pipe_kwargs)
697
 
698
  generated_image = result.images[0]
699
 
700
- # Post-processing
701
  if enable_color_matching and has_detected_faces:
702
  print("Applying enhanced face-aware color matching...")
703
  try:
@@ -721,7 +406,11 @@ class RetroArtConverter:
721
  except Exception as e:
722
  print(f"Color matching failed: {e}")
723
 
 
 
 
 
724
  return generated_image
725
 
726
 
727
- print("[OK] Generator class ready")
 
1
  """
2
  Generation logic for Pixagram AI Pixel Art Generator
3
+ UPDATED VERSION with simplified InstantID face preservation
4
  """
5
  import torch
6
  import numpy as np
 
15
  )
16
  from utils import (
17
  sanitize_text, enhanced_color_match, color_match, create_face_mask,
18
+ draw_kps, get_demographic_description, calculate_optimal_size, enhance_face_crop,
19
+ safe_image_size, ensure_int
20
  )
21
  from models import (
22
+ load_face_analysis, load_depth_detector, load_controlnets,
23
+ load_sdxl_pipeline, load_lora, setup_compel,
24
  setup_scheduler, optimize_pipeline, load_caption_model, set_clip_skip
25
  )
26
 
27
 
28
  class RetroArtConverter:
29
+ """Main class for retro art generation with InstantID face preservation"""
30
 
31
  def __init__(self):
32
  self.device = device
 
35
  'custom_checkpoint': False,
36
  'lora': False,
37
  'instantid': False,
38
+ 'zoe_depth': False
 
39
  }
40
 
41
  # Initialize face analysis
 
45
  self.zoe_depth, zoe_success = load_depth_detector()
46
  self.models_loaded['zoe_depth'] = zoe_success
47
 
48
+ # Load ControlNets - ALWAYS as list for InstantID pipeline
49
+ controlnet_instantid, controlnet_depth = load_controlnets()
50
+ controlnets = [controlnet_instantid, controlnet_depth]
51
+ self.models_loaded['instantid'] = True
 
52
 
53
+ print("Initializing InstantID pipeline with Face + Depth ControlNets")
 
 
 
 
 
 
 
 
 
 
 
 
54
 
55
+ # Load SDXL pipeline with InstantID (handles IP-Adapter internally)
56
  self.pipe, checkpoint_success = load_sdxl_pipeline(controlnets)
57
  self.models_loaded['custom_checkpoint'] = checkpoint_success
58
 
 
60
  lora_success = load_lora(self.pipe)
61
  self.models_loaded['lora'] = lora_success
62
 
 
 
 
 
 
 
 
 
 
63
  # Setup Compel
64
  self.compel, self.use_compel = setup_compel(self.pipe)
65
 
 
78
  print(" [OK] Using GIT for detailed captions")
79
  elif self.caption_model_type == "blip":
80
  print(" [OK] Using BLIP for standard captions")
 
 
 
81
 
82
  # Set CLIP skip
83
  set_clip_skip(self.pipe)
84
 
 
 
 
 
85
  # Print model status
86
  self._print_status()
87
 
88
+ print(" [OK] Model initialization complete with InstantID!")
89
 
90
  def _print_status(self):
91
  """Print model loading status"""
 
93
  for model, loaded in self.models_loaded.items():
94
  status = "[OK] LOADED" if loaded else "[FALLBACK/DISABLED]"
95
  print(f"{model}: {status}")
96
+ print("InstantID Pipeline: [OK] Active with built-in IP-Adapter")
97
  print("===================\n")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98
 
99
  def get_depth_map(self, image):
100
+ """Generate depth map using Zoe Depth"""
101
+ if self.zoe_depth is not None:
102
+ try:
103
+ if image.mode != 'RGB':
104
+ image = image.convert('RGB')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
105
 
106
+ # Use safe helpers for type safety
107
+ orig_width, orig_height = safe_image_size(image)
108
 
109
+ # FIXED: Use multiples of 64 (not 32)
110
+ target_width = ensure_int((orig_width // 64) * 64)
111
+ target_height = ensure_int((orig_height // 64) * 64)
112
 
113
+ target_width = ensure_int(max(64, target_width))
114
+ target_height = ensure_int(max(64, target_height))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
115
 
116
+ # Create an explicit tuple of standard ints
117
+ size_for_depth = (target_width, target_height)
 
 
 
 
 
 
 
 
118
 
119
+ # Always resize using the explicit int tuple
120
+ image_for_depth = image.resize(size_for_depth, Image.LANCZOS)
121
 
122
+ # Generate depth map
123
+ depth_image = self.zoe_depth(image_for_depth, detect_resolution=512, image_resolution=1024)
 
124
 
125
+ # Resize to match original if needed
126
+ if (depth_image.width, depth_image.height) != (orig_width, orig_height):
127
+ depth_image = depth_image.resize((orig_width, orig_height), Image.LANCZOS)
 
 
 
 
 
 
 
128
 
129
+ # Convert to RGB if needed
130
+ if depth_image.mode != 'RGB':
131
+ depth_image = depth_image.convert('RGB')
132
 
133
+ return depth_image, np.array(depth_image)
 
 
134
 
135
+ except Exception as e:
136
+ print(f"Depth map generation failed: {e}")
137
+ import traceback
138
+ traceback.print_exc()
139
+ return None, None
140
+ else:
141
+ print(" Zoe Depth not available")
142
+ return None, None
143
+
144
+ def generate_caption(self, image):
145
+ """Generate caption for image using loaded caption model"""
146
+ if not self.caption_enabled or self.caption_model is None:
147
+ return None
148
+
149
+ try:
150
+ if self.caption_model_type == 'git':
151
+ # GIT model
152
+ pixel_values = self.caption_processor(images=image, return_tensors="pt").pixel_values
153
+ pixel_values = pixel_values.to(device=self.device, dtype=self.dtype)
154
+
155
+ generated_ids = self.caption_model.generate(
156
+ pixel_values=pixel_values,
157
+ max_length=CAPTION_CONFIG['max_length']
158
+ )
159
+ caption = self.caption_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
160
 
161
+ elif self.caption_model_type == 'blip':
162
+ # BLIP model
163
+ inputs = self.caption_processor(image, return_tensors="pt").to(self.device, self.dtype)
164
+ out = self.caption_model.generate(**inputs, max_new_tokens=CAPTION_CONFIG['max_length'])
165
+ caption = self.caption_processor.decode(out[0], skip_special_tokens=True)
166
+
167
+ else:
168
+ return None
169
+
170
+ return sanitize_text(caption)
171
 
 
 
172
  except Exception as e:
173
  print(f"Caption generation failed: {e}")
174
  return None
 
176
  def generate_retro_art(
177
  self,
178
  input_image,
179
+ prompt,
180
+ negative_prompt="",
181
  num_inference_steps=12,
182
+ guidance_scale=1.3,
183
+ depth_control_scale=0.75,
184
  identity_control_scale=0.85,
185
  lora_scale=1.0,
186
+ identity_preservation=1.2,
187
+ strength=0.50,
188
  enable_color_matching=False,
189
  consistency_mode=True,
190
  seed=-1
191
  ):
192
+ """
193
+ Generate retro art with InstantID face preservation.
 
 
 
194
 
195
+ UPDATED: Simplified face embedding handling using InstantID pipeline.
196
+ """
 
 
197
 
198
+ # Validate and adjust parameters if consistency mode is enabled
199
  if consistency_mode:
200
+ # Ensure guidance scale is in optimal range for LCM
201
+ if guidance_scale < 1.0:
202
+ guidance_scale = 1.0
203
+ elif guidance_scale > 1.8:
204
+ guidance_scale = 1.8
205
+
206
+ # Ensure identity preservation and lora scale balance
207
+ if identity_preservation > 1.5 and lora_scale > 1.2:
208
+ lora_scale = min(lora_scale, 1.0)
209
+
210
+ # Ensure strength is reasonable
211
+ if strength < 0.3:
212
+ strength = 0.3
213
+ elif strength > 0.8:
214
+ strength = 0.8
215
 
216
+ # Calculate optimal size
217
+ orig_width, orig_height = safe_image_size(input_image)
218
+ optimal_width, optimal_height = calculate_optimal_size(orig_width, orig_height)
219
 
220
+ # Resize image
221
+ resized_image = input_image.resize((optimal_width, optimal_height), Image.LANCZOS)
222
 
223
  # Generate depth map
224
+ print("Generating depth map...")
225
+ depth_image, depth_array = self.get_depth_map(resized_image)
226
+
227
+ if depth_image is None:
228
+ raise RuntimeError("Failed to generate depth map")
229
 
230
+ # Detect faces
231
+ print("Detecting faces...")
232
+ has_detected_faces = False
233
  face_kps_image = None
234
  face_embeddings = None
 
 
235
  face_bbox_original = None
236
 
237
+ if self.face_detection_enabled and self.face_app is not None:
238
+ try:
239
+ faces = self.face_app.get(np.array(resized_image))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
240
 
241
+ if len(faces) > 0:
242
+ has_detected_faces = True
243
+ face = faces[0]
244
+
245
+ # Draw keypoints
246
+ face_kps_image = draw_kps(resized_image, face.kps)
247
+
248
+ # Get face embeddings (512D vector from InsightFace)
249
+ face_embeddings = face.embedding
250
+
251
+ # Get face bounding box for color matching
252
+ face_bbox_original = face.bbox
253
+
254
+ print(f" [OK] Face detected")
255
+ print(f" - Embedding shape: {face_embeddings.shape}")
256
+ print(f" - Keypoints: {face.kps.shape}")
257
+ print(f" - Bbox: {face_bbox_original}")
258
+
259
+ # Check for adaptive parameter adjustment
260
+ face_area = (face.bbox[2] - face.bbox[0]) * (face.bbox[3] - face.bbox[1])
261
+ det_score = face.det_score if hasattr(face, 'det_score') else 1.0
262
+
263
+ # Apply adaptive adjustments
264
+ if face_area < ADAPTIVE_THRESHOLDS['small_face_size']:
265
+ print(" [ADAPTIVE] Small face detected - boosting preservation")
266
+ identity_preservation = max(identity_preservation, ADAPTIVE_PARAMS['small_face']['identity_preservation'])
267
+ identity_control_scale = max(identity_control_scale, ADAPTIVE_PARAMS['small_face']['identity_control_scale'])
268
+
269
+ elif det_score < ADAPTIVE_THRESHOLDS['low_confidence']:
270
+ print(" [ADAPTIVE] Low confidence - increasing identity weight")
271
+ identity_preservation = max(identity_preservation, ADAPTIVE_PARAMS['low_confidence']['identity_preservation'])
272
+ identity_control_scale = max(identity_control_scale, ADAPTIVE_PARAMS['low_confidence']['identity_control_scale'])
273
+
274
+ else:
275
+ print(" No faces detected in image")
276
+
277
+ except Exception as e:
278
+ print(f"Face detection error: {e}")
279
+ has_detected_faces = False
280
+
281
+ # Enhance prompt with trigger word
282
+ if TRIGGER_WORD not in prompt.lower():
283
+ prompt = f"{TRIGGER_WORD}, {prompt}"
284
 
285
  # Set LORA scale
286
  if hasattr(self.pipe, 'set_adapters') and self.models_loaded['lora']:
 
335
  if hasattr(self.pipe, 'text_encoder'):
336
  pipe_kwargs["clip_skip"] = 2
337
 
338
+ # ========================================
339
+ # SIMPLIFIED: Configure ControlNets + IP-Adapter
340
+ # ========================================
341
+ if has_detected_faces and face_kps_image is not None:
342
+ print("Using InstantID (keypoints + embeddings) + Depth ControlNets")
343
 
344
+ # Control images: [face keypoints, depth map]
345
+ pipe_kwargs["control_image"] = [face_kps_image, depth_image]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
346
 
347
+ # Conditioning scales: [identity, depth]
348
+ pipe_kwargs["controlnet_conditioning_scale"] = [
349
+ identity_control_scale,
350
+ depth_control_scale
351
+ ]
 
 
 
 
352
 
353
+ # CRITICAL: Pass face embeddings for IP-Adapter
354
+ # The InstantID pipeline handles the Resampler internally!
355
+ if face_embeddings is not None:
356
+ print(f"Adding face embeddings for IP-Adapter...")
357
+
358
+ # Just pass the embeddings - pipeline does everything!
359
+ pipe_kwargs["image_embeds"] = face_embeddings # numpy array (512,)
360
+
361
+ # Control IP-Adapter strength
362
+ boosted_scale = identity_preservation * IDENTITY_BOOST_MULTIPLIER
363
+ pipe_kwargs["ip_adapter_scale"] = boosted_scale
364
+
365
+ print(f" - Face embeddings shape: {face_embeddings.shape}")
366
+ print(f" - IP-Adapter scale: {boosted_scale:.2f}")
367
+ print(f" [OK] Face embeddings configured")
368
+ else:
369
+ print(" [WARNING] No face embeddings - using keypoints only")
370
 
371
  else:
372
+ print("No faces detected - using Depth ControlNet only")
373
+
374
+ # Use depth for both ControlNet slots (identity scale = 0)
375
+ pipe_kwargs["control_image"] = [depth_image, depth_image]
376
+ pipe_kwargs["controlnet_conditioning_scale"] = [0.0, depth_control_scale]
377
 
378
  # Generate
379
  print(f"Generating with LCM: Steps={num_inference_steps}, CFG={guidance_scale}, Strength={strength}")
380
+ print(f"ControlNet scales - Identity: {identity_control_scale}, Depth: {depth_control_scale}")
381
  result = self.pipe(**pipe_kwargs)
382
 
383
  generated_image = result.images[0]
384
 
385
+ # Post-processing: Color matching
386
  if enable_color_matching and has_detected_faces:
387
  print("Applying enhanced face-aware color matching...")
388
  try:
 
406
  except Exception as e:
407
  print(f"Color matching failed: {e}")
408
 
409
+ # Memory cleanup
410
+ if torch.cuda.is_available():
411
+ torch.cuda.empty_cache()
412
+
413
  return generated_image
414
 
415
 
416
+ print("[OK] Generator class ready with InstantID pipeline")