Spaces:
Runtime error
Runtime error
Update generator.py
Browse files- 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,
|
| 21 |
-
load_sdxl_pipeline, load_lora,
|
| 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 |
-
|
| 49 |
-
|
| 50 |
-
self.
|
| 51 |
-
self.models_loaded['instantid'] = instantid_success
|
| 52 |
|
| 53 |
-
|
| 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 |
-
|
| 150 |
-
|
| 151 |
-
|
| 152 |
-
|
| 153 |
-
|
| 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 |
-
#
|
| 226 |
-
|
| 227 |
|
| 228 |
-
#
|
| 229 |
-
|
| 230 |
-
|
| 231 |
|
| 232 |
-
|
| 233 |
-
|
| 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 |
-
|
| 382 |
-
|
| 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 |
-
|
|
|
|
| 393 |
|
| 394 |
-
|
| 395 |
-
|
| 396 |
-
inputs = self.caption_processor(images=image, return_tensors="pt").to(self.device, self.dtype)
|
| 397 |
|
| 398 |
-
|
| 399 |
-
|
| 400 |
-
|
| 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 |
-
|
|
|
|
|
|
|
| 410 |
|
| 411 |
-
|
| 412 |
-
# BLIP base processing
|
| 413 |
-
inputs = self.caption_processor(image, return_tensors="pt").to(self.device, self.dtype)
|
| 414 |
|
| 415 |
-
|
| 416 |
-
|
| 417 |
-
|
| 418 |
-
|
| 419 |
-
|
| 420 |
-
|
| 421 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 422 |
|
| 423 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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
|
| 435 |
-
negative_prompt="
|
| 436 |
num_inference_steps=12,
|
| 437 |
-
guidance_scale=1.
|
| 438 |
-
depth_control_scale=0.
|
| 439 |
identity_control_scale=0.85,
|
| 440 |
lora_scale=1.0,
|
| 441 |
-
identity_preservation=
|
| 442 |
-
strength=0.
|
| 443 |
enable_color_matching=False,
|
| 444 |
consistency_mode=True,
|
| 445 |
seed=-1
|
| 446 |
):
|
| 447 |
-
"""
|
| 448 |
-
|
| 449 |
-
# Sanitize text inputs
|
| 450 |
-
prompt = sanitize_text(prompt)
|
| 451 |
-
negative_prompt = sanitize_text(negative_prompt)
|
| 452 |
|
| 453 |
-
|
| 454 |
-
|
| 455 |
-
negative_prompt = ""
|
| 456 |
-
# **FIX 3 END**
|
| 457 |
|
| 458 |
-
#
|
| 459 |
if consistency_mode:
|
| 460 |
-
|
| 461 |
-
|
| 462 |
-
|
| 463 |
-
|
| 464 |
-
|
| 465 |
-
|
| 466 |
-
|
| 467 |
-
|
| 468 |
-
|
| 469 |
-
|
| 470 |
-
|
| 471 |
-
|
| 472 |
-
|
|
|
|
|
|
|
| 473 |
|
| 474 |
-
|
| 475 |
-
|
| 476 |
-
|
| 477 |
|
| 478 |
-
# Resize
|
| 479 |
-
resized_image = input_image.resize((
|
| 480 |
|
| 481 |
# Generate depth map
|
| 482 |
-
print("Generating
|
| 483 |
-
depth_image = self.get_depth_map(resized_image)
|
| 484 |
-
|
| 485 |
-
|
|
|
|
| 486 |
|
| 487 |
-
#
|
| 488 |
-
|
|
|
|
| 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
|
| 496 |
-
|
| 497 |
-
|
| 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 |
-
|
| 560 |
-
|
| 561 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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 |
-
#
|
| 617 |
-
|
| 618 |
-
|
| 619 |
-
|
| 620 |
-
|
| 621 |
|
| 622 |
-
|
| 623 |
-
pipe_kwargs["
|
| 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 |
-
|
| 675 |
-
|
| 676 |
-
|
| 677 |
-
|
| 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 |
-
|
| 685 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 686 |
|
| 687 |
else:
|
| 688 |
-
print("
|
| 689 |
-
|
| 690 |
-
|
| 691 |
-
|
|
|
|
| 692 |
|
| 693 |
# Generate
|
| 694 |
print(f"Generating with LCM: Steps={num_inference_steps}, CFG={guidance_scale}, Strength={strength}")
|
| 695 |
-
print(f"
|
| 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")
|