local_inpainter.py

"""

Local inpainting implementation - COMPATIBLE VERSION WITH JIT SUPPORT

Maintains full backward compatibility while adding proper JIT model support

"""
import os
import sys
import json
import numpy as np
import cv2
from typing import Optional, List, Tuple, Dict, Any
import logging
import traceback
import re
import hashlib
import urllib.request
from pathlib import Path
import threading
import time

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Check if we're running in a frozen environment
IS_FROZEN = getattr(sys, 'frozen', False)
if IS_FROZEN:
    MEIPASS = sys._MEIPASS
    os.environ['TORCH_HOME'] = MEIPASS
    os.environ['TRANSFORMERS_CACHE'] = os.path.join(MEIPASS, 'transformers')
    os.environ['HF_HOME'] = os.path.join(MEIPASS, 'huggingface')
    logger.info(f"Running in frozen environment: {MEIPASS}")

# Environment variables for ONNX
ONNX_CACHE_DIR = os.environ.get('ONNX_CACHE_DIR', 'models')
AUTO_CONVERT_TO_ONNX = os.environ.get('AUTO_CONVERT_TO_ONNX', 'false').lower() == 'true'
SKIP_ONNX_FOR_CKPT = os.environ.get('SKIP_ONNX_FOR_CKPT', 'true').lower() == 'true'
FORCE_ONNX_REBUILD = os.environ.get('FORCE_ONNX_REBUILD', 'false').lower() == 'true'
CACHE_DIR = os.environ.get('MODEL_CACHE_DIR', os.path.expanduser('~/.cache/inpainting'))

# Modified import handling for frozen environment
TORCH_AVAILABLE = False
torch = None
nn = None
F = None
BaseModel = object

try:
    import onnxruntime_extensions
    ONNX_EXTENSIONS_AVAILABLE = True
except ImportError:
    ONNX_EXTENSIONS_AVAILABLE = False
    logger.info("ONNX Runtime Extensions not available - FFT models won't work in ONNX")

if IS_FROZEN:
    # In frozen environment, try harder to import
    try:
        import torch
        import torch.nn as nn
        import torch.nn.functional as F
        TORCH_AVAILABLE = True
        BaseModel = nn.Module
        logger.info("✓ PyTorch loaded in frozen environment")
    except Exception as e:
        logger.error(f"PyTorch not available in frozen environment: {e}")
        logger.error("❌ Inpainting disabled - PyTorch is required")
else:
    # Normal environment
    try:
        import torch
        import torch.nn as nn
        import torch.nn.functional as F
        TORCH_AVAILABLE = True
        BaseModel = nn.Module
    except ImportError:
        TORCH_AVAILABLE = False
        logger.error("PyTorch not available - inpainting disabled")

# Configure ORT memory behavior before importing
try:
    os.environ.setdefault('ORT_DISABLE_MEMORY_ARENA', '1')
except Exception:
    pass
# ONNX Runtime - usually works well in frozen environments
ONNX_AVAILABLE = False
try:
    import onnx
    import onnxruntime as ort
    ONNX_AVAILABLE = True
    logger.info("✓ ONNX Runtime available")
except ImportError:
    ONNX_AVAILABLE = False
    logger.warning("ONNX Runtime not available")

# Bubble detector - optional
BUBBLE_DETECTOR_AVAILABLE = False
try:
    from bubble_detector import BubbleDetector
    BUBBLE_DETECTOR_AVAILABLE = True
    logger.info("✓ Bubble detector available")
except ImportError:
    logger.info("Bubble detector not available - basic inpainting will be used")


# JIT Model URLs (for automatic download)
LAMA_JIT_MODELS = {
    'lama': {
        'url': 'https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt',
        'md5': 'e3aa4aaa15225a33ec84f9f4bc47e500',
        'name': 'BigLama'
    },
    'anime': {
        'url': 'https://github.com/Sanster/models/releases/download/AnimeMangaInpainting/anime-manga-big-lama.pt',
        'md5': '29f284f36a0a510bcacf39ecf4c4d54f',
        'name': 'Anime-Manga BigLama'
    },
    'lama_official': {
        'url': 'https://github.com/Sanster/models/releases/download/lama/lama.pt',
        'md5': '4b1a1de53b7a74e0ff9dd622834e8e1e',
        'name': 'LaMa Official'
    },
    'aot': {
        'url': 'https://huggingface.co/ogkalu/aot-inpainting-jit/resolve/main/aot_traced.pt',
        'md5': '5ecdac562c1d56267468fc4fbf80db27',
        'name': 'AOT GAN'
    },
    'aot_onnx': {
        'url': 'https://huggingface.co/ogkalu/aot-inpainting/resolve/main/aot.onnx',
        'md5': 'ffd39ed8e2a275869d3b49180d030f0d8b8b9c2c20ed0e099ecd207201f0eada',
        'name': 'AOT ONNX (Fast)',
        'is_onnx': True
    },
    'lama_onnx': {
        'url': 'https://huggingface.co/Carve/LaMa-ONNX/resolve/main/lama_fp32.onnx',
        'md5': None,  # Add MD5 if you want to verify
        'name': 'LaMa ONNX (Carve)',
        'is_onnx': True  # Flag to indicate this is ONNX, not JIT
    },
    'anime_onnx': {
        'url': 'https://huggingface.co/ogkalu/lama-manga-onnx-dynamic/resolve/main/lama-manga-dynamic.onnx',
        'md5': 'de31ffa5ba26916b8ea35319f6c12151ff9654d4261bccf0583a69bb095315f9',
        'name': 'Anime/Manga ONNX (Dynamic)',
        'is_onnx': True  # Flag to indicate this is ONNX
    }
}


def norm_img(img: np.ndarray) -> np.ndarray:
    """Normalize image to [0, 1] range"""
    if img.dtype == np.uint8:
        return img.astype(np.float32) / 255.0
    return img


def get_cache_path_by_url(url: str) -> str:
    """Get cache path for a model URL"""
    os.makedirs(CACHE_DIR, exist_ok=True)
    filename = os.path.basename(url)
    return os.path.join(CACHE_DIR, filename)


def download_model(url: str, md5: str = None) -> str:
    """Download model if not cached"""
    cache_path = get_cache_path_by_url(url)
    
    if os.path.exists(cache_path):
        logger.info(f"✅ Model already cached: {cache_path}")
        return cache_path
    
    logger.info(f"📥 Downloading model from {url}")
    
    try:
        urllib.request.urlretrieve(url, cache_path)
        logger.info(f"✅ Model downloaded to: {cache_path}")
        return cache_path
    except Exception as e:
        logger.error(f"❌ Download failed: {e}")
        if os.path.exists(cache_path):
            os.remove(cache_path)
        raise


class FFCInpaintModel(BaseModel):  # Use BaseModel instead of nn.Module
    """FFC model for LaMa inpainting - for checkpoint compatibility"""
    
    def __init__(self):
        if not TORCH_AVAILABLE:
            # Initialize as a simple object when PyTorch is not available
            super().__init__()
            logger.warning("PyTorch not available - FFCInpaintModel initialized as placeholder")
            self._pytorch_available = False
            return
            
        # Additional safety check for nn being None
        if nn is None:
            super().__init__()
            logger.error("Neural network modules not available - FFCInpaintModel disabled")
            self._pytorch_available = False
            return
            
        super().__init__()
        self._pytorch_available = True
        
        try:
            # Encoder
            self.model_1_ffc_convl2l = nn.Conv2d(4, 64, 7, padding=3)
            self.model_1_bn_l = nn.BatchNorm2d(64)
            
            self.model_2_ffc_convl2l = nn.Conv2d(64, 128, 3, padding=1)
            self.model_2_bn_l = nn.BatchNorm2d(128)
            
            self.model_3_ffc_convl2l = nn.Conv2d(128, 256, 3, padding=1)
            self.model_3_bn_l = nn.BatchNorm2d(256)
            
            self.model_4_ffc_convl2l = nn.Conv2d(256, 128, 3, padding=1)
            self.model_4_ffc_convl2g = nn.Conv2d(256, 384, 3, padding=1)
            self.model_4_bn_l = nn.BatchNorm2d(128)
            self.model_4_bn_g = nn.BatchNorm2d(384)
            
            # FFC blocks
            for i in range(5, 23):
                for conv_type in ['conv1', 'conv2']:
                    setattr(self, f'model_{i}_{conv_type}_ffc_convl2l', nn.Conv2d(128, 128, 3, padding=1))
                    setattr(self, f'model_{i}_{conv_type}_ffc_convl2g', nn.Conv2d(128, 384, 3, padding=1))
                    setattr(self, f'model_{i}_{conv_type}_ffc_convg2l', nn.Conv2d(384, 128, 3, padding=1))
                    setattr(self, f'model_{i}_{conv_type}_ffc_convg2g_conv1_0', nn.Conv2d(384, 192, 1))
                    setattr(self, f'model_{i}_{conv_type}_ffc_convg2g_conv1_1', nn.BatchNorm2d(192))
                    setattr(self, f'model_{i}_{conv_type}_ffc_convg2g_fu_conv_layer', nn.Conv2d(384, 384, 1))
                    setattr(self, f'model_{i}_{conv_type}_ffc_convg2g_fu_bn', nn.BatchNorm2d(384))
                    setattr(self, f'model_{i}_{conv_type}_ffc_convg2g_conv2', nn.Conv2d(192, 384, 1))
                    setattr(self, f'model_{i}_{conv_type}_bn_l', nn.BatchNorm2d(128))
                    setattr(self, f'model_{i}_{conv_type}_bn_g', nn.BatchNorm2d(384))
            
            # Decoder
            self.model_24 = nn.Conv2d(512, 256, 3, padding=1)
            self.model_25 = nn.BatchNorm2d(256)
            
            self.model_27 = nn.Conv2d(256, 128, 3, padding=1)
            self.model_28 = nn.BatchNorm2d(128)
            
            self.model_30 = nn.Conv2d(128, 64, 3, padding=1)
            self.model_31 = nn.BatchNorm2d(64)
            
            self.model_34 = nn.Conv2d(64, 3, 7, padding=3)
            
            # Activation functions
            self.relu = nn.ReLU(inplace=True)
            self.tanh = nn.Tanh()
            
            logger.info("FFCInpaintModel initialized successfully")
            
        except Exception as e:
            logger.error(f"Failed to initialize FFCInpaintModel: {e}")
            self._pytorch_available = False
            raise
    
    def forward(self, image, mask):
        if not self._pytorch_available:
            logger.error("PyTorch not available for forward pass")
            raise RuntimeError("PyTorch not available for forward pass")
            
        if not TORCH_AVAILABLE or torch is None:
            logger.error("PyTorch not available for forward pass")
            raise RuntimeError("PyTorch not available for forward pass")
            
        try:
            x = torch.cat([image, mask], dim=1)
            
            x = self.relu(self.model_1_bn_l(self.model_1_ffc_convl2l(x)))
            x = self.relu(self.model_2_bn_l(self.model_2_ffc_convl2l(x)))
            x = self.relu(self.model_3_bn_l(self.model_3_ffc_convl2l(x)))
            
            x_l = self.relu(self.model_4_bn_l(self.model_4_ffc_convl2l(x)))
            x_g = self.relu(self.model_4_bn_g(self.model_4_ffc_convl2g(x)))
            
            for i in range(5, 23):
                identity_l, identity_g = x_l, x_g
                x_l, x_g = self._ffc_block(x_l, x_g, i, 'conv1')
                x_l, x_g = self._ffc_block(x_l, x_g, i, 'conv2')
                x_l = x_l + identity_l
                x_g = x_g + identity_g
            
            x = torch.cat([x_l, x_g], dim=1)
            x = self.relu(self.model_25(self.model_24(x)))
            x = self.relu(self.model_28(self.model_27(x)))
            x = self.relu(self.model_31(self.model_30(x)))
            x = self.tanh(self.model_34(x))
            
            mask_3ch = mask.repeat(1, 3, 1, 1)
            return x * mask_3ch + image * (1 - mask_3ch)
            
        except Exception as e:
            logger.error(f"Forward pass failed: {e}")
            raise RuntimeError(f"Forward pass failed: {e}")
    
    def _ffc_block(self, x_l, x_g, idx, conv_type):
        if not self._pytorch_available:
            raise RuntimeError("PyTorch not available for FFC block")
            
        if not TORCH_AVAILABLE:
            raise RuntimeError("PyTorch not available for FFC block")
            
        try:
            convl2l = getattr(self, f'model_{idx}_{conv_type}_ffc_convl2l')
            convl2g = getattr(self, f'model_{idx}_{conv_type}_ffc_convl2g')
            convg2l = getattr(self, f'model_{idx}_{conv_type}_ffc_convg2l')
            convg2g_conv1 = getattr(self, f'model_{idx}_{conv_type}_ffc_convg2g_conv1_0')
            convg2g_bn1 = getattr(self, f'model_{idx}_{conv_type}_ffc_convg2g_conv1_1')
            fu_conv = getattr(self, f'model_{idx}_{conv_type}_ffc_convg2g_fu_conv_layer')
            fu_bn = getattr(self, f'model_{idx}_{conv_type}_ffc_convg2g_fu_bn')
            convg2g_conv2 = getattr(self, f'model_{idx}_{conv_type}_ffc_convg2g_conv2')
            bn_l = getattr(self, f'model_{idx}_{conv_type}_bn_l')
            bn_g = getattr(self, f'model_{idx}_{conv_type}_bn_g')
            
            out_xl = convl2l(x_l) + convg2l(x_g)
            out_xg = convl2g(x_l) + convg2g_conv2(self.relu(convg2g_bn1(convg2g_conv1(x_g)))) + self.relu(fu_bn(fu_conv(x_g)))
            
            return self.relu(bn_l(out_xl)), self.relu(bn_g(out_xg))
            
        except Exception as e:
            logger.error(f"FFC block failed: {e}")
            raise RuntimeError(f"FFC block failed: {e}")


class LocalInpainter:
    """Local inpainter with full backward compatibility"""
    
    # MAINTAIN ORIGINAL SUPPORTED_METHODS for compatibility
    SUPPORTED_METHODS = {
        'lama': ('LaMa Inpainting', FFCInpaintModel),
        'mat': ('MAT Inpainting', FFCInpaintModel),
        'aot': ('AOT GAN Inpainting', FFCInpaintModel),
        'aot_onnx': ('AOT ONNX (Fast)', FFCInpaintModel),
        'sd': ('Stable Diffusion Inpainting', FFCInpaintModel),
        'anime': ('Anime/Manga Inpainting', FFCInpaintModel),
        'anime_onnx': ('Anime ONNX (Fast)', FFCInpaintModel),
        'lama_official': ('Official LaMa', FFCInpaintModel),
    }
    
    def __init__(self, config_path="config.json"):
        # Set thread limits early if environment indicates single-threaded mode
        try:
            if os.environ.get('OMP_NUM_THREADS') == '1':
                # Already in single-threaded mode, ensure it's applied to this process
                # Check if torch is available at module level before trying to use it
                if TORCH_AVAILABLE and torch is not None:
                    try:
                        torch.set_num_threads(1)
                    except (RuntimeError, AttributeError):
                        pass
                try:
                    import cv2
                    cv2.setNumThreads(1)
                except (ImportError, AttributeError):
                    pass
        except Exception:
            pass
        
        self.config_path = config_path
        self.config = self._load_config()
        self.model = None
        self.model_loaded = False
        self.current_method = None
        self.use_opencv_fallback = False  # FORCED DISABLED - No OpenCV fallback allowed
        self.onnx_session = None
        self.use_onnx = False
        self.is_jit_model = False
        self.pad_mod = 8
        
        # Default tiling settings - OFF by default for most models
        self.tiling_enabled = False
        self.tile_size = 512
        self.tile_overlap = 64
        
        # ONNX-specific settings
        self.onnx_model_loaded = False
        self.onnx_input_size = None  # Will be detected from model
        
        # Quantization diagnostics flags
        self.onnx_quantize_applied = False
        self.torch_quantize_applied = False
        
        # Bubble detection
        self.bubble_detector = None
        self.bubble_model_loaded = False
        
        # Create directories
        os.makedirs(ONNX_CACHE_DIR, exist_ok=True)
        os.makedirs(CACHE_DIR, exist_ok=True)
        logger.info(f"📁 ONNX cache directory: {ONNX_CACHE_DIR}")
        logger.info(f"   Contents: {os.listdir(ONNX_CACHE_DIR) if os.path.exists(ONNX_CACHE_DIR) else 'Directory does not exist'}")

        
        # Check GPU availability safely
        self.use_gpu = False
        self.device = None
        
        if TORCH_AVAILABLE and torch is not None:
            try:
                self.use_gpu = torch.cuda.is_available()
                self.device = torch.device('cuda' if self.use_gpu else 'cpu')
                if self.use_gpu:
                    logger.info(f"🚀 GPU: {torch.cuda.get_device_name(0)}")
                else:
                    logger.info("💻 Using CPU")
            except AttributeError:
                # torch module exists but doesn't have cuda attribute
                self.use_gpu = False
                self.device = None
                logger.info("⚠️ PyTorch incomplete - inpainting disabled")
        else:
            logger.info("⚠️ PyTorch not available - inpainting disabled")
        
        # Quantization/precision toggle (off by default)
        try:
            adv_cfg = self.config.get('manga_settings', {}).get('advanced', {}) if isinstance(self.config, dict) else {}
            # Track singleton mode from settings for thread limiting (deprecated - kept for compatibility)
            self.singleton_mode = bool(adv_cfg.get('use_singleton_models', True))
            env_quant = os.environ.get('MODEL_QUANTIZE', 'false').lower() == 'true'
            self.quantize_enabled = bool(env_quant or adv_cfg.get('quantize_models', False))
            # ONNX quantization is now strictly opt-in (config or env), decoupled from general quantize_models
            self.onnx_quantize_enabled = bool(
                adv_cfg.get('onnx_quantize', os.environ.get('ONNX_QUANTIZE', 'false').lower() == 'true')
            )
            self.torch_precision = str(adv_cfg.get('torch_precision', os.environ.get('TORCH_PRECISION', 'auto'))).lower()
            logger.info(f"Quantization: {'ENABLED' if self.quantize_enabled else 'disabled'} for Local Inpainter; onnx_quantize={'on' if self.onnx_quantize_enabled else 'off'}; torch_precision={self.torch_precision}")
            self.int8_enabled = bool(
                adv_cfg.get('int8_quantize', False)
                or adv_cfg.get('quantize_int8', False)
                or os.environ.get('TORCH_INT8', 'false').lower() == 'true'
                or self.torch_precision in ('int8', 'int8_dynamic')
            )
            logger.info(
                f"Quantization: {'ENABLED' if self.quantize_enabled else 'disabled'} for Local Inpainter; "
                f"onnx_quantize={'on' if self.onnx_quantize_enabled else 'off'}; "
                f"torch_precision={self.torch_precision}; int8={'on' if self.int8_enabled else 'off'}"
            )
        except Exception:
            self.quantize_enabled = False
            self.onnx_quantize_enabled = False
            self.torch_precision = 'auto'
            self.int8_enabled = False
        
        # HD strategy defaults (mirror of comic-translate behavior)
        try:
            adv_cfg = self.config.get('manga_settings', {}).get('advanced', {}) if isinstance(self.config, dict) else {}
        except Exception:
            adv_cfg = {}
        try:
            self.hd_strategy = str(os.environ.get('HD_STRATEGY', adv_cfg.get('hd_strategy', 'resize'))).lower()
        except Exception:
            self.hd_strategy = 'resize'
        try:
            self.hd_resize_limit = int(os.environ.get('HD_RESIZE_LIMIT', adv_cfg.get('hd_strategy_resize_limit', 1536)))
        except Exception:
            self.hd_resize_limit = 1536
        try:
            self.hd_crop_margin = int(os.environ.get('HD_CROP_MARGIN', adv_cfg.get('hd_strategy_crop_margin', 16)))
        except Exception:
            self.hd_crop_margin = 16
        try:
            self.hd_crop_trigger_size = int(os.environ.get('HD_CROP_TRIGGER', adv_cfg.get('hd_strategy_crop_trigger_size', 1024)))
        except Exception:
            self.hd_crop_trigger_size = 1024
        logger.info(f"HD strategy: {self.hd_strategy} (resize_limit={self.hd_resize_limit}, crop_margin={self.hd_crop_margin}, crop_trigger={self.hd_crop_trigger_size})")
        
        # Stop flag support
        self.stop_flag = None
        self._stopped = False
        self.log_callback = None
        
        # Initialize bubble detector if available
        if BUBBLE_DETECTOR_AVAILABLE:
            try:
                self.bubble_detector = BubbleDetector()
                logger.info("🗨️ Bubble detection available")
            except:
                self.bubble_detector = None
                logger.info("🗨️ Bubble detection not available")
    
    def _load_config(self):
        try:
            if self.config_path and os.path.exists(self.config_path):
                with open(self.config_path, 'r', encoding='utf-8') as f:
                    content = f.read().strip()
                    if not content:
                        return {}
                    try:
                        return json.loads(content)
                    except json.JSONDecodeError:
                        # Likely a concurrent write; retry once after a short delay
                        try:
                            import time
                            time.sleep(0.05)
                            with open(self.config_path, 'r', encoding='utf-8') as f2:
                                return json.load(f2)
                        except Exception:
                            return {}
        except Exception:
            return {}
        return {}
    
    def _save_config(self):
        # Don't save if config is empty (prevents purging)
        if not getattr(self, 'config', None):
            return
        try:
            # Load existing (best-effort)
            full_config = {}
            if self.config_path and os.path.exists(self.config_path):
                try:
                    with open(self.config_path, 'r', encoding='utf-8') as f:
                        full_config = json.load(f)
                except Exception as read_err:
                    logger.debug(f"Config read during save failed (non-critical): {read_err}")
                    full_config = {}
            # Update
            full_config.update(self.config)
            # Atomic write: write to temp then replace
            tmp_path = (self.config_path or 'config.json') + '.tmp'
            with open(tmp_path, 'w', encoding='utf-8') as f:
                json.dump(full_config, f, indent=2, ensure_ascii=False)
            try:
                os.replace(tmp_path, self.config_path or 'config.json')
            except Exception as replace_err:
                logger.debug(f"Config atomic replace failed, trying direct write: {replace_err}")
                # Fallback to direct write
                with open(self.config_path or 'config.json', 'w', encoding='utf-8') as f:
                    json.dump(full_config, f, indent=2, ensure_ascii=False)
        except Exception as save_err:
            # Never crash on config save, but log for debugging
            logger.debug(f"Config save failed (non-critical): {save_err}")
            pass
    
    def set_stop_flag(self, stop_flag):
        """Set the stop flag for checking interruptions"""
        self.stop_flag = stop_flag
        self._stopped = False
    
    def set_log_callback(self, log_callback):
        """Set log callback for GUI integration"""
        self.log_callback = log_callback
    
    def _check_stop(self) -> bool:
        """Check if stop has been requested"""
        if self._stopped:
            return True
        if self.stop_flag and self.stop_flag.is_set():
            self._stopped = True
            return True
        # Check global manga translator cancellation
        try:
            from manga_translator import MangaTranslator
            if MangaTranslator.is_globally_cancelled():
                self._stopped = True
                return True
        except Exception:
            pass
        return False
    
    def _log(self, message: str, level: str = "info"):
        """Log message with stop suppression"""
        # Suppress logs when stopped (allow only essential stop confirmation messages)
        if self._check_stop():
            essential_stop_keywords = [
                "⏹️ Translation stopped by user",
                "⏹️ Inpainting stopped",
                "cleanup", "🧹"
            ]
            if not any(keyword in message for keyword in essential_stop_keywords):
                return
        
        if self.log_callback:
            self.log_callback(message, level)
        else:
            logger.info(message) if level == 'info' else getattr(logger, level, logger.info)(message)
    
    def reset_stop_flags(self):
        """Reset stop flags when starting new processing"""
        self._stopped = False

    def convert_to_onnx(self, model_path: str, method: str) -> Optional[str]:
        """Convert a PyTorch model to ONNX format with FFT handling via custom operators"""
        if not ONNX_AVAILABLE:
            logger.warning("ONNX not available, skipping conversion")
            return None
        
        try:
            # Generate ONNX path
            model_name = os.path.basename(model_path).replace('.pt', '')
            onnx_path = os.path.join(ONNX_CACHE_DIR, f"{model_name}_{method}.onnx")
            
            # Check if ONNX already exists
            if os.path.exists(onnx_path) and not FORCE_ONNX_REBUILD:
                logger.info(f"✅ ONNX model already exists: {onnx_path}")
                return onnx_path
            
            logger.info(f"🔄 Converting {method} model to ONNX...")
            
            # The model should already be loaded at this point
            if not self.model_loaded or self.current_method != method:
                logger.error("Model not loaded for ONNX conversion")
                return None
            
            # Create dummy inputs
            dummy_image = torch.randn(1, 3, 512, 512).to(self.device)
            dummy_mask = torch.randn(1, 1, 512, 512).to(self.device)
            
            # For FFT models, we can't convert directly
            fft_models = ['lama', 'anime', 'lama_official']
            if method in fft_models:
                logger.warning(f"⚠️ {method.upper()} uses FFT operations that cannot be exported")
                return None  # Just return None, don't suggest Carve
            
            # Standard export for non-FFT models
            try:
                torch.onnx.export(
                    self.model,
                    (dummy_image, dummy_mask),
                    onnx_path,
                    export_params=True,
                    opset_version=13,
                    do_constant_folding=True,
                    input_names=['image', 'mask'],
                    output_names=['output'],
                    dynamic_axes={
                        'image': {0: 'batch', 2: 'height', 3: 'width'},
                        'mask': {0: 'batch', 2: 'height', 3: 'width'},
                        'output': {0: 'batch', 2: 'height', 3: 'width'}
                    }
                )
                logger.info(f"✅ ONNX model saved to: {onnx_path}")
                return onnx_path
                
            except torch.onnx.errors.UnsupportedOperatorError as e:
                logger.error(f"❌ Unsupported operator: {e}")
                return None
            
        except Exception as e:
            logger.error(f"❌ ONNX conversion failed: {e}")
            logger.error(traceback.format_exc())
            return None

    def load_onnx_model(self, onnx_path: str) -> bool:
        """Load an ONNX model with custom operator support"""
        if not ONNX_AVAILABLE:
            logger.error("ONNX Runtime not available")
            return False
        
        # Check if this exact ONNX model is already loaded
        if (self.onnx_session is not None and 
            hasattr(self, 'current_onnx_path') and 
            self.current_onnx_path == onnx_path):
            logger.debug(f"✅ ONNX model already loaded: {onnx_path}")
            return True
        
        try:
            # Don't log here if we already logged in load_model
            logger.debug(f"📦 ONNX Runtime loading: {onnx_path}")
            
            # Store the path for later checking
            self.current_onnx_path = onnx_path

            # Check if this is a Carve model (fixed 512x512)
            is_carve_model = "lama_fp32" in onnx_path or "carve" in onnx_path.lower()
            if is_carve_model:
                logger.info("📦 Detected Carve ONNX model (fixed 512x512 input)")
                self.onnx_fixed_size = (512, 512)
            else:
                self.onnx_fixed_size = None
            
            # Standard ONNX loading: prefer CUDA if available; otherwise CPU. Do NOT use DML.
            try:
                avail = ort.get_available_providers() if ONNX_AVAILABLE else []
            except Exception:
                avail = []
            if 'CUDAExecutionProvider' in avail:
                providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
            else:
                providers = ['CPUExecutionProvider']
            session_path = onnx_path
            try:
                fname_lower = os.path.basename(onnx_path).lower()
            except Exception:
                fname_lower = str(onnx_path).lower()

            # Device-aware policy for LaMa-type ONNX (Carve or contains 'lama')
            is_lama_model = is_carve_model or ('lama' in fname_lower)
            if is_lama_model:
                base = os.path.splitext(onnx_path)[0]
                if self.use_gpu:
                    # Prefer FP16 on CUDA
                    fp16_path = base + '.fp16.onnx'
                    if (not os.path.exists(fp16_path)) or FORCE_ONNX_REBUILD:
                        try:
                            import onnx as _onnx
                            try:
                                from onnxruntime_tools.transformers.float16 import convert_float_to_float16 as _to_fp16
                            except Exception:
                                try:
                                    from onnxconverter_common import float16
                                    def _to_fp16(m, keep_io_types=True):
                                        return float16.convert_float_to_float16(m, keep_io_types=keep_io_types)
                                except Exception:
                                    _to_fp16 = None
                            if _to_fp16 is not None:
                                m = _onnx.load(onnx_path)
                                m_fp16 = _to_fp16(m, keep_io_types=True)
                                _onnx.save(m_fp16, fp16_path)
                                logger.info(f"✅ Generated FP16 ONNX for LaMa: {fp16_path}")
                        except Exception as e:
                            logger.warning(f"FP16 conversion for LaMa failed: {e}")
                    if os.path.exists(fp16_path):
                        session_path = fp16_path
                else:
                    # CPU path for LaMa: quantize only if enabled, and MatMul-only to avoid artifacts
                    if self.onnx_quantize_enabled:
                        try:
                            from onnxruntime.quantization import quantize_dynamic, QuantType
                            quant_path = base + '.matmul.int8.onnx'
                            if (not os.path.exists(quant_path)) or FORCE_ONNX_REBUILD:
                                logger.info("🔻 LaMa: Quantizing ONNX weights to INT8 (dynamic, ops=['MatMul'])...")
                                quantize_dynamic(
                                    model_input=onnx_path,
                                    model_output=quant_path,
                                    weight_type=QuantType.QInt8,
                                    op_types_to_quantize=['MatMul']
                                )
                            self.onnx_quantize_applied = True
                            # Validate dynamic quant result
                            try:
                                import onnx as _onnx
                                _m_q = _onnx.load(quant_path)
                                _onnx.checker.check_model(_m_q)
                            except Exception as _qchk:
                                logger.warning(f"LaMa dynamic quant model invalid; deleting and falling back: {_qchk}")
                                try:
                                    os.remove(quant_path)
                                except Exception:
                                    pass
                                quant_path = None
                        except Exception as dy_err:
                            logger.warning(f"LaMa dynamic quantization failed: {dy_err}")
                            quant_path = None
                        # Fallback: static QDQ MatMul-only with zero data reader
                        if quant_path is None:
                            try:
                                import onnx as _onnx
                                from onnxruntime.quantization import (
                                    CalibrationDataReader, quantize_static,
                                    QuantFormat, QuantType, CalibrationMethod
                                )
                                m = _onnx.load(onnx_path)
                                shapes = {}
                                for inp in m.graph.input:
                                    dims = []
                                    for d in inp.type.tensor_type.shape.dim:
                                        dims.append(d.dim_value if d.dim_value > 0 else 1)
                                    shapes[inp.name] = dims
                                class _ZeroReader(CalibrationDataReader):
                                    def __init__(self, shapes):
                                        self.shapes = shapes
                                        self.done = False
                                    def get_next(self):
                                        if self.done:
                                            return None
                                        feed = {}
                                        for name, s in self.shapes.items():
                                            ss = list(s)
                                            if len(ss) == 4:
                                                if ss[2] <= 1: ss[2] = 512
                                                if ss[3] <= 1: ss[3] = 512
                                                if ss[1] <= 1 and 'mask' not in name.lower():
                                                    ss[1] = 3
                                            feed[name] = np.zeros(ss, dtype=np.float32)
                                        self.done = True
                                        return feed
                                dr = _ZeroReader(shapes)
                                quant_path = base + '.matmul.int8.onnx'
                                quantize_static(
                                    model_input=onnx_path,
                                    model_output=quant_path,
                                    calibration_data_reader=dr,
                                    quant_format=QuantFormat.QDQ,
                                    activation_type=QuantType.QUInt8,
                                    weight_type=QuantType.QInt8,
                                    per_channel=False,
                                    calibrate_method=CalibrationMethod.MinMax,
                                    op_types_to_quantize=['MatMul']
                                )
                                # Validate
                                try:
                                    _m_q = _onnx.load(quant_path)
                                    _onnx.checker.check_model(_m_q)
                                except Exception as _qchk2:
                                    logger.warning(f"LaMa static MatMul-only quant model invalid; deleting: {_qchk2}")
                                    try:
                                        os.remove(quant_path)
                                    except Exception:
                                        pass
                                    quant_path = None
                                else:
                                    logger.info(f"✅ Generated MatMul-only INT8 ONNX for LaMa: {quant_path}")
                                    self.onnx_quantize_applied = True
                            except Exception as st_err:
                                logger.warning(f"LaMa static MatMul-only quantization failed: {st_err}")
                                quant_path = None
                        # Use the quantized model if valid
                        if quant_path and os.path.exists(quant_path):
                            session_path = quant_path
                            logger.info(f"✅ Using LaMa quantized ONNX model: {quant_path}")
                    # If quantization not enabled or failed, session_path remains onnx_path (FP32)

            # Optional dynamic/static quantization for other models (opt-in)
            if (not is_lama_model) and self.onnx_quantize_enabled:
                base = os.path.splitext(onnx_path)[0]
                fname = os.path.basename(onnx_path).lower()
                is_aot = 'aot' in fname
                # For AOT: ignore any MatMul-only file and prefer Conv+MatMul
                if is_aot:
                    try:
                        ignored_matmul = base + ".matmul.int8.onnx"
                        if os.path.exists(ignored_matmul):
                            logger.info(f"⏭️ Ignoring MatMul-only quantized file for AOT: {ignored_matmul}")
                    except Exception:
                        pass
                # Choose target quant file and ops
                if is_aot:
                    quant_path = base + ".int8.onnx"
                    ops_to_quant = ['MatMul']
                    # Use MatMul-only for safer quantization across models
                    ops_for_static = ['MatMul']
                    # Try to simplify AOT graph prior to quantization
                    quant_input_path = onnx_path
                    try:
                        import onnx as _onnx
                        try:
                            from onnxsim import simplify as _onnx_simplify
                            _model = _onnx.load(onnx_path)
                            _sim_model, _check = _onnx_simplify(_model)
                            if _check:
                                sim_path = base + ".sim.onnx"
                                _onnx.save(_sim_model, sim_path)
                                quant_input_path = sim_path
                                logger.info(f"🧰 Simplified AOT ONNX before quantization: {sim_path}")
                        except Exception as _sim_err:
                            logger.info(f"AOT simplification skipped: {_sim_err}")
                        # No ONNX shape inference; keep original graph structure
                        # Ensure opset >= 13 for QDQ (axis attribute on DequantizeLinear)
                        try:
                            _m_tmp = _onnx.load(quant_input_path)
                            _opset = max([op.version for op in _m_tmp.opset_import]) if _m_tmp.opset_import else 11
                            if _opset < 13:
                                from onnx import version_converter as _vc
                                _m13 = _vc.convert_version(_m_tmp, 13)
                                up_path = base + ".op13.onnx"
                                _onnx.save(_m13, up_path)
                                quant_input_path = up_path
                                logger.info(f"🧰 Upgraded ONNX opset to 13 before QDQ quantization: {up_path}")
                        except Exception as _operr:
                            logger.info(f"Opset upgrade skipped: {_operr}")
                    except Exception:
                        quant_input_path = onnx_path
                else:
                    quant_path = base + ".matmul.int8.onnx"
                    ops_to_quant = ['MatMul']
                    ops_for_static = ops_to_quant
                    quant_input_path = onnx_path
                # Perform quantization if needed
                if not os.path.exists(quant_path) or FORCE_ONNX_REBUILD:
                    if is_aot:
                        # Directly perform static QDQ quantization for MatMul only (avoid Conv activations)
                        try:
                            import onnx as _onnx
                            from onnxruntime.quantization import CalibrationDataReader, quantize_static, QuantFormat, QuantType, CalibrationMethod
                            _model = _onnx.load(quant_input_path)
                            # Build input shapes from the model graph
                            input_shapes = {}
                            for inp in _model.graph.input:
                                dims = []
                                for d in inp.type.tensor_type.shape.dim:
                                    if d.dim_value > 0:
                                        dims.append(d.dim_value)
                                    else:
                                        # default fallback dimension
                                        dims.append(1)
                                input_shapes[inp.name] = dims
                            class _ZeroDataReader(CalibrationDataReader):
                                def __init__(self, input_shapes):
                                    self._shapes = input_shapes
                                    self._provided = False
                                def get_next(self):
                                    if self._provided:
                                        return None
                                    feed = {}
                                    for name, shape in self._shapes.items():
                                        # Ensure reasonable default spatial size
                                        s = list(shape)
                                        if len(s) == 4:
                                            if s[2] <= 1:
                                                s[2] = 512
                                            if s[3] <= 1:
                                                s[3] = 512
                                            # channel fallback
                                            if s[1] <= 1 and 'mask' not in name.lower():
                                                s[1] = 3
                                        feed[name] = (np.zeros(s, dtype=np.float32))
                                    self._provided = True
                                    return feed
                            dr = _ZeroDataReader(input_shapes)
                            quantize_static(
                                model_input=quant_input_path,
                                model_output=quant_path,
                                calibration_data_reader=dr,
                                quant_format=QuantFormat.QDQ,
                                activation_type=QuantType.QUInt8,
                                weight_type=QuantType.QInt8,
                                per_channel=True,
                                calibrate_method=CalibrationMethod.MinMax,
                                op_types_to_quantize=ops_for_static
                            )
                            # Validate quantized model to catch structural errors early
                            try:
                                _m_q = _onnx.load(quant_path)
                                _onnx.checker.check_model(_m_q)
                            except Exception as _qchk:
                                logger.warning(f"Quantized AOT model validation failed: {_qchk}")
                                # Remove broken quantized file to force fallback
                                try:
                                    os.remove(quant_path)
                                except Exception:
                                    pass
                            else:
                                logger.info(f"✅ Static INT8 quantization produced: {quant_path}")
                        except Exception as st_err:
                            logger.warning(f"Static ONNX quantization failed: {st_err}")
                    else:
                        # First attempt: dynamic quantization (MatMul)
                        try:
                            from onnxruntime.quantization import quantize_dynamic, QuantType
                            logger.info("🔻 Quantizing ONNX inpainting model weights to INT8 (dynamic, ops=['MatMul'])...")
                            quantize_dynamic(
                                model_input=quant_input_path,
                                model_output=quant_path,
                                weight_type=QuantType.QInt8,
                                op_types_to_quantize=['MatMul']
                            )
                        except Exception as dy_err:
                            logger.warning(f"Dynamic ONNX quantization failed: {dy_err}; attempting static quantization...")
                            # Fallback: static quantization with a zero data reader
                            try:
                                import onnx as _onnx
                                from onnxruntime.quantization import CalibrationDataReader, quantize_static, QuantFormat, QuantType, CalibrationMethod
                                _model = _onnx.load(quant_input_path)
                                # Build input shapes from the model graph
                                input_shapes = {}
                                for inp in _model.graph.input:
                                    dims = []
                                    for d in inp.type.tensor_type.shape.dim:
                                        if d.dim_value > 0:
                                            dims.append(d.dim_value)
                                        else:
                                            # default fallback dimension
                                            dims.append(1)
                                    input_shapes[inp.name] = dims
                                class _ZeroDataReader(CalibrationDataReader):
                                    def __init__(self, input_shapes):
                                        self._shapes = input_shapes
                                        self._provided = False
                                    def get_next(self):
                                        if self._provided:
                                            return None
                                        feed = {}
                                        for name, shape in self._shapes.items():
                                            # Ensure reasonable default spatial size
                                            s = list(shape)
                                            if len(s) == 4:
                                                if s[2] <= 1:
                                                    s[2] = 512
                                                if s[3] <= 1:
                                                    s[3] = 512
                                                # channel fallback
                                                if s[1] <= 1 and 'mask' not in name.lower():
                                                    s[1] = 3
                                            feed[name] = (np.zeros(s, dtype=np.float32))
                                        self._provided = True
                                        return feed
                                dr = _ZeroDataReader(input_shapes)
                                quantize_static(
                                    model_input=quant_input_path,
                                    model_output=quant_path,
                                    calibration_data_reader=dr,
                                    quant_format=QuantFormat.QDQ,
                                    activation_type=QuantType.QUInt8,
                                    weight_type=QuantType.QInt8,
                                    per_channel=True,
                                    calibrate_method=CalibrationMethod.MinMax,
                                    op_types_to_quantize=ops_for_static
                                )
                                # Validate quantized model to catch structural errors early
                                try:
                                    _m_q = _onnx.load(quant_path)
                                    _onnx.checker.check_model(_m_q)
                                except Exception as _qchk:
                                    logger.warning(f"Quantized AOT model validation failed: {_qchk}")
                                    # Remove broken quantized file to force fallback
                                    try:
                                        os.remove(quant_path)
                                    except Exception:
                                        pass
                                else:
                                    logger.info(f"✅ Static INT8 quantization produced: {quant_path}")
                            except Exception as st_err:
                                logger.warning(f"Static ONNX quantization failed: {st_err}")
                # Prefer the quantized file if it now exists
                if os.path.exists(quant_path):
                    # Validate existing quantized model before using it
                    try:
                        import onnx as _onnx
                        _m_q = _onnx.load(quant_path)
                        _onnx.checker.check_model(_m_q)
                    except Exception as _qchk:
                        logger.warning(f"Existing quantized ONNX invalid; deleting and falling back: {_qchk}")
                        try:
                            os.remove(quant_path)
                        except Exception:
                            pass
                    else:
                        session_path = quant_path
                        logger.info(f"✅ Using quantized ONNX model: {quant_path}")
                else:
                    logger.warning("ONNX quantization not applied: quantized file not created")
            # Use conservative ORT memory options to reduce RAM growth
            so = ort.SessionOptions()
            try:
                so.enable_mem_pattern = False
                so.enable_cpu_mem_arena = False
            except Exception:
                pass
            # Enable optimal performance settings (let ONNX use all CPU cores)
            try:
                # Use all available CPU threads for best performance
                # ONNX Runtime will automatically use optimal thread count
                so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED
            except Exception:
                pass
            # Try to create an inference session, with graceful fallbacks
            try:
                self.onnx_session = ort.InferenceSession(session_path, sess_options=so, providers=providers)
            except Exception as e:
                err = str(e)
                logger.warning(f"ONNX session creation failed for {session_path}: {err}")
                # If quantized path failed due to unsupported ops or invalid graph, remove it and retry unquantized
                if session_path != onnx_path and ('ConvInteger' in err or 'NOT_IMPLEMENTED' in err or 'INVALID_ARGUMENT' in err):
                    try:
                        if os.path.exists(session_path):
                            os.remove(session_path)
                            logger.info(f"🧹 Deleted invalid quantized model: {session_path}")
                    except Exception:
                        pass
                    try:
                        logger.info("Retrying with unquantized ONNX model...")
                        self.onnx_session = ort.InferenceSession(onnx_path, sess_options=so, providers=providers)
                        session_path = onnx_path
                    except Exception as e2:
                        logger.warning(f"Unquantized ONNX session failed with current providers: {e2}")
                        # As a last resort, try CPU-only
                        try:
                            logger.info("Retrying ONNX on CPUExecutionProvider only...")
                            self.onnx_session = ort.InferenceSession(onnx_path, sess_options=so, providers=['CPUExecutionProvider'])
                            session_path = onnx_path
                            providers = ['CPUExecutionProvider']
                        except Exception as e3:
                            logger.error(f"Failed to create ONNX session on CPU: {e3}")
                            raise
                else:
                    # If we weren't quantized but failed on CUDA, try CPU-only once
                    if self.use_gpu and 'NOT_IMPLEMENTED' in err:
                        try:
                            logger.info("Retrying ONNX on CPUExecutionProvider only...")
                            self.onnx_session = ort.InferenceSession(session_path, sess_options=so, providers=['CPUExecutionProvider'])
                            providers = ['CPUExecutionProvider']
                        except Exception as e4:
                            logger.error(f"Failed to create ONNX session on CPU: {e4}")
                            raise
            
            # Get input/output names
            if self.onnx_session is None:
                raise RuntimeError("ONNX session was not created")
            self.onnx_input_names = [i.name for i in self.onnx_session.get_inputs()]
            self.onnx_output_names = [o.name for o in self.onnx_session.get_outputs()]
            
            # Check input shapes to detect fixed-size models
            input_shape = self.onnx_session.get_inputs()[0].shape
            if len(input_shape) == 4 and input_shape[2] == 512 and input_shape[3] == 512:
                self.onnx_fixed_size = (512, 512)
                logger.info(f"  Model expects fixed size: 512x512")
            
            # Log success with I/O info in a single line
            logger.debug(f"✅ ONNX session created - Inputs: {self.onnx_input_names}, Outputs: {self.onnx_output_names}")
            
            self.use_onnx = True
            # CRITICAL: Set model_loaded flag when ONNX session is successfully created
            # This ensures preloaded spares are recognized as valid loaded instances
            self.model_loaded = True
            return True
            
        except Exception as e:
            logger.error(f"❌ Failed to load ONNX: {e}")
            import traceback
            logger.debug(f"ONNX load traceback: {traceback.format_exc()}")
            self.use_onnx = False
            self.model_loaded = False
            return False
        
    def _convert_checkpoint_key(self, key):
        """Convert checkpoint key format to model format"""
        # model.24.weight -> model_24.weight
        if re.match(r'^model\.(\d+)\.(weight|bias|running_mean|running_var)$', key):
            return re.sub(r'model\.(\d+)\.', r'model_\1.', key)
        
        # model.5.conv1.ffc.weight -> model_5_conv1_ffc.weight
        if key.startswith('model.'):
            parts = key.split('.')
            if parts[-1] in ['weight', 'bias', 'running_mean', 'running_var']:
                return '_'.join(parts[:-1]).replace('model_', 'model_') + '.' + parts[-1]
        
        return key.replace('.', '_')
    
    def _load_weights_with_mapping(self, model, state_dict):
        """Load weights with proper mapping"""
        model_dict = model.state_dict()
        
        logger.info(f"📊 Model expects {len(model_dict)} weights")
        logger.info(f"📊 Checkpoint has {len(state_dict)} weights")
        
        # Filter out num_batches_tracked
        actual_weights = {k: v for k, v in state_dict.items() if 'num_batches_tracked' not in k}
        logger.info(f"   Actual weights: {len(actual_weights)}")
        
        mapped = {}
        unmapped_ckpt = []
        unmapped_model = list(model_dict.keys())
        
        # Map checkpoint weights
        for ckpt_key, ckpt_val in actual_weights.items():
            success = False
            converted_key = self._convert_checkpoint_key(ckpt_key)
            
            if converted_key in model_dict:
                target_shape = model_dict[converted_key].shape
                
                if target_shape == ckpt_val.shape:
                    mapped[converted_key] = ckpt_val
                    success = True
                elif len(ckpt_val.shape) == 4 and len(target_shape) == 4:
                    # 4D permute for decoder convs
                    permuted = ckpt_val.permute(1, 0, 2, 3)
                    if target_shape == permuted.shape:
                        mapped[converted_key] = permuted
                        logger.info(f"   ✅ Permuted: {ckpt_key}")
                        success = True
                elif len(ckpt_val.shape) == 2 and len(target_shape) == 2:
                    # 2D transpose
                    transposed = ckpt_val.transpose(0, 1)
                    if target_shape == transposed.shape:
                        mapped[converted_key] = transposed
                        success = True
                
                if success and converted_key in unmapped_model:
                    unmapped_model.remove(converted_key)
            
            if not success:
                unmapped_ckpt.append(ckpt_key)
        
        # Try fallback mapping for unmapped
        if unmapped_ckpt:
            logger.info(f"   🔧 Fallback mapping for {len(unmapped_ckpt)} weights...")
            for ckpt_key in unmapped_ckpt[:]:
                ckpt_val = actual_weights[ckpt_key]
                for model_key in unmapped_model[:]:
                    if model_dict[model_key].shape == ckpt_val.shape:
                        if ('weight' in ckpt_key and 'weight' in model_key) or \
                           ('bias' in ckpt_key and 'bias' in model_key):
                            mapped[model_key] = ckpt_val
                            unmapped_model.remove(model_key)
                            unmapped_ckpt.remove(ckpt_key)
                            logger.info(f"   ✅ Mapped: {ckpt_key} -> {model_key}")
                            break
        
        # Initialize missing weights
        complete_dict = model_dict.copy()
        complete_dict.update(mapped)
        
        for key in unmapped_model:
            param = complete_dict[key]
            if 'weight' in key:
                if 'conv' in key.lower():
                    nn.init.kaiming_normal_(param, mode='fan_out', nonlinearity='relu')
                else:
                    nn.init.xavier_uniform_(param)
            elif 'bias' in key:
                nn.init.zeros_(param)
            elif 'running_mean' in key:
                nn.init.zeros_(param)
            elif 'running_var' in key:
                nn.init.ones_(param)
        
        # Report
        logger.info(f"✅ Mapped {len(actual_weights) - len(unmapped_ckpt)}/{len(actual_weights)} checkpoint weights")
        logger.info(f"   Filled {len(mapped)}/{len(model_dict)} model positions")
        
        if unmapped_model:
            pct = (len(unmapped_model) / len(model_dict)) * 100
            logger.info(f"   ⚠️ Initialized {len(unmapped_model)} missing weights ({pct:.1f}%)")
            if pct > 20:
                logger.warning("   ⚠️ May produce artifacts - checkpoint is incomplete")
                logger.warning("   💡 Consider downloading JIT model for better quality:")
                logger.warning(f"      inpainter.download_jit_model('{self.current_method or 'lama'}')")
        
        model.load_state_dict(complete_dict, strict=True)
        return True
    
    def download_jit_model(self, method: str) -> str:
        """Download JIT model for a method"""
        if method in LAMA_JIT_MODELS:
            model_info = LAMA_JIT_MODELS[method]
            logger.info(f"📥 Downloading {model_info['name']}...")
            
            try:
                model_path = download_model(model_info['url'], model_info['md5'])
                return model_path
            except Exception as e:
                logger.error(f"Failed to download {method}: {e}")
        else:
            logger.warning(f"No JIT model available for {method}")
        
        return None
    
    def load_model(self, method, model_path, force_reload=False):
        """Load model - supports both JIT and checkpoint files with ONNX conversion"""
        try:
            if not TORCH_AVAILABLE:
                logger.warning("PyTorch not available in this build")
                logger.info("Inpainting features will be disabled - this is normal for lightweight builds")
                logger.info("The application will continue to work without local inpainting")
                self.model_loaded = False
                return False
            
            # Additional safety check for torch being None
            if torch is None or nn is None:
                logger.warning("PyTorch modules not properly loaded")
                logger.info("Inpainting features will be disabled - this is normal for lightweight builds")
                self.model_loaded = False
                return False
            
            # Check if model path changed - but only if we had a previous path saved
            current_saved_path = self.config.get(f'{method}_model_path', '')
            if current_saved_path and current_saved_path != model_path:
                logger.info(f"📍 Model path changed for {method}")
                logger.info(f"   Old: {current_saved_path}")
                logger.info(f"   New: {model_path}")
                force_reload = True
            
            if not os.path.exists(model_path):
                # Try to auto-download JIT model if path doesn't exist
                logger.warning(f"Model not found: {model_path}")
                logger.info("Attempting to download JIT model...")
                
                try:
                    jit_path = self.download_jit_model(method)
                    if jit_path and os.path.exists(jit_path):
                        model_path = jit_path
                        logger.info(f"Using downloaded JIT model: {jit_path}")
                    else:
                        logger.error(f"Model not found and download failed: {model_path}")
                        logger.info("Inpainting will be unavailable for this session")
                        return False
                except Exception as download_error:
                    logger.error(f"Download failed: {download_error}")
                    logger.info("Inpainting will be unavailable for this session")
                    return False
            
            # Check if already loaded in THIS instance
            if self.model_loaded and self.current_method == method and not force_reload:
                # Additional check for ONNX - make sure the session exists
                if self.use_onnx and self.onnx_session is not None:
                    logger.debug(f"✅ {method.upper()} ONNX already loaded (skipping reload)")
                    return True
                elif not self.use_onnx and self.model is not None:
                    logger.debug(f"✅ {method.upper()} already loaded (skipping reload)")
                    return True
                else:
                    # Model claims to be loaded but objects are missing - force reload
                    logger.warning(f"⚠️ Model claims loaded but session/model object is None - forcing reload")
                    force_reload = True
                    self.model_loaded = False
            
            # Clear previous model if force reload
            if force_reload:
                logger.info(f"🔄 Force reloading {method} model...")
                self.model = None
                self.onnx_session = None
                self.model_loaded = False
                self.is_jit_model = False
                # Only log loading message when actually loading
                logger.info(f"📥 Loading {method} from {model_path}")
            elif self.model_loaded and self.current_method != method:
                # If we have a model loaded but it's a different method, clear it
                logger.info(f"🔄 Switching from {self.current_method} to {method}")
                self.model = None
                self.onnx_session = None
                self.model_loaded = False
                self.is_jit_model = False
                # Only log loading message when actually loading
                logger.info(f"📥 Loading {method} from {model_path}")
            elif not self.model_loaded:
                # Only log when we're actually going to load
                logger.info(f"📥 Loading {method} from {model_path}")
            # else: model is loaded and current, no logging needed

            # Normalize path and enforce expected extension for certain methods
            try:
                _ext = os.path.splitext(model_path)[1].lower()
                _method_lower = str(method).lower()
                # For explicit ONNX methods, ensure we use a .onnx path
                if _method_lower in ("lama_onnx", "anime_onnx", "aot_onnx") and _ext != ".onnx":
                    # If the file exists, try to detect if it's actually an ONNX model and correct the extension
                    if os.path.exists(model_path) and ONNX_AVAILABLE:
                        try:
                            import onnx as _onnx
                            _ = _onnx.load(model_path)  # will raise if not ONNX
                            # Build a corrected path under the ONNX cache dir
                            base_name = os.path.splitext(os.path.basename(model_path))[0]
                            if base_name.endswith('.pt'):
                                base_name = base_name[:-3]
                            corrected_path = os.path.join(ONNX_CACHE_DIR, base_name + ".onnx")
                            # Avoid overwriting a valid file with an invalid one
                            if model_path != corrected_path:
                                try:
                                    import shutil as _shutil
                                    _shutil.copy2(model_path, corrected_path)
                                    model_path = corrected_path
                                    logger.info(f"🔧 Corrected ONNX model extension/path: {model_path}")
                                except Exception as _cp_e:
                                    # As a fallback, try in-place rename to .onnx
                                    try:
                                        in_place = os.path.splitext(model_path)[0] + ".onnx"
                                        os.replace(model_path, in_place)
                                        model_path = in_place
                                        logger.info(f"🔧 Renamed ONNX model to: {model_path}")
                                    except Exception:
                                        logger.warning(f"Could not correct ONNX extension automatically: {_cp_e}")
                        except Exception:
                            # Not an ONNX file; leave as-is
                            pass
                    # If the path doesn't exist or still wrong, prefer the known ONNX download for this method
                    if (not os.path.exists(model_path)) or (os.path.splitext(model_path)[1].lower() != ".onnx"):
                        try:
                            # Download the appropriate ONNX model based on the method
                            if _method_lower == "anime_onnx":
                                _dl = self.download_jit_model("anime_onnx")
                            elif _method_lower == "aot_onnx":
                                _dl = self.download_jit_model("aot_onnx")
                            else:
                                _dl = self.download_jit_model("lama_onnx")
                            if _dl and os.path.exists(_dl):
                                model_path = _dl
                                logger.info(f"🔧 Using downloaded {_method_lower.upper()} model: {model_path}")
                        except Exception:
                            pass
            except Exception:
                pass

            # Check file signature to detect ONNX files (even with wrong extension)
            # or check file extension
            ext = model_path.lower().split('.')[-1]
            is_onnx = False
            
            # Check by file signature
            try:
                with open(model_path, 'rb') as f:
                    file_header = f.read(8)
                if file_header.startswith(b'\x08'):
                    is_onnx = True
                    logger.debug("📦 Detected ONNX file signature")
            except Exception:
                pass
            
            # Check by extension
            if ext == 'onnx':
                is_onnx = True
            
            # Handle ONNX files
            if is_onnx:
                # Note: load_onnx_model will handle its own logging
                try:
                    onnx_load_result = self.load_onnx_model(model_path)
                    if onnx_load_result:
                        # CRITICAL: Set model_loaded flag FIRST before any other operations
                        # This ensures concurrent threads see the correct state immediately
                        self.model_loaded = True
                        self.use_onnx = True
                        self.is_jit_model = False
                        # Ensure aot_onnx is properly set as current method
                        if 'aot' in method.lower():
                            self.current_method = 'aot_onnx'
                        else:
                            self.current_method = method
                        # Save with BOTH key formats for compatibility (non-critical - do last)
                        try:
                            self.config[f'{method}_model_path'] = model_path
                            self.config[f'manga_{method}_model_path'] = model_path
                            self._save_config()
                        except Exception as cfg_err:
                            logger.debug(f"Config save after ONNX load failed (non-critical): {cfg_err}")
                        logger.info(f"✅ {method.upper()} ONNX loaded with method: {self.current_method}")
                        # Double-check model_loaded flag is still set
                        if not self.model_loaded:
                            logger.error("❌ CRITICAL: model_loaded flag was unset after successful ONNX load!")
                            self.model_loaded = True
                        return True
                    else:
                        logger.error("Failed to load ONNX model - load_onnx_model returned False")
                        self.model_loaded = False
                        return False
                except Exception as onnx_err:
                    logger.error(f"Exception during ONNX model loading: {onnx_err}")
                    import traceback
                    logger.debug(traceback.format_exc())
                    self.model_loaded = False
                    return False
            
            # Check if it's a JIT model (.pt) or checkpoint (.ckpt/.pth)
            if model_path.endswith('.pt'):
                try:
                    # Try loading as JIT/TorchScript
                    logger.info("📦 Attempting to load as JIT model...")
                    self.model = torch.jit.load(model_path, map_location=self.device or 'cpu')
                    self.model.eval()
                    
                    if self.use_gpu and self.device:
                        try:
                            self.model = self.model.to(self.device)
                        except Exception as gpu_error:
                            logger.warning(f"Could not move model to GPU: {gpu_error}")
                            logger.info("Using CPU instead")
                    
                    self.is_jit_model = True
                    self.model_loaded = True
                    self.current_method = method
                    logger.info("✅ JIT model loaded successfully!")
                    time.sleep(0.1)  # Brief pause for stability
                    logger.debug("💤 JIT model loading pausing briefly for stability")
                    
                    # Optional FP16 precision on GPU to reduce VRAM
                    if self.quantize_enabled and self.use_gpu:
                        try:
                            if self.torch_precision in ('fp16', 'auto'):
                                self.model = self.model.half()
                                logger.info("🔻 Applied FP16 precision to inpainting model (GPU)")
                            else:
                                logger.info("Torch precision set to fp32; skipping half()")
                        except Exception as _e:
                            logger.warning(f"Could not switch inpainting model precision: {_e}")
                    
                    # Optional INT8 dynamic quantization for CPU TorchScript (best-effort)
                    if (self.int8_enabled or (self.quantize_enabled and not self.use_gpu and self.torch_precision in ('auto', 'int8'))) and not self.use_gpu:
                        try:
                            applied = False
                            # Try TorchScript dynamic quantization API (older PyTorch)
                            try:
                                from torch.quantization import quantize_dynamic_jit  # type: ignore
                                self.model = quantize_dynamic_jit(self.model, {"aten::linear"}, dtype=torch.qint8)  # type: ignore
                                applied = True
                            except Exception:
                                pass
                            # Try eager-style dynamic quantization on the scripted module (may no-op)
                            if not applied:
                                try:
                                    import torch.ao.quantization as tq  # type: ignore
                                    self.model = tq.quantize_dynamic(self.model, {nn.Linear}, dtype=torch.qint8)  # type: ignore
                                    applied = True
                                except Exception:
                                    pass
                            # Always try to optimize TorchScript for inference
                            try:
                                self.model = torch.jit.optimize_for_inference(self.model)  # type: ignore
                            except Exception:
                                pass
                            if applied:
                                logger.info("🔻 Applied INT8 dynamic quantization to JIT inpainting model (CPU)")
                                self.torch_quantize_applied = True
                            else:
                                logger.info("ℹ️ INT8 dynamic quantization not applied (unsupported for this JIT graph); using FP32 CPU")
                        except Exception as _qe:
                            logger.warning(f"INT8 quantization skipped: {_qe}")
                    
                    # Save with BOTH key formats for compatibility
                    self.config[f'{method}_model_path'] = model_path
                    self.config[f'manga_{method}_model_path'] = model_path
                    self._save_config()
                    
                    # ONNX CONVERSION (optionally in background)
                    if AUTO_CONVERT_TO_ONNX and self.model_loaded:
                        def _convert_and_switch():
                            try:
                                onnx_path = self.convert_to_onnx(model_path, method)
                                if onnx_path and self.load_onnx_model(onnx_path):
                                    logger.info("🚀 Using ONNX model for inference")
                                else:
                                    logger.info("📦 Using PyTorch JIT model for inference")
                            except Exception as onnx_error:
                                logger.warning(f"ONNX conversion failed: {onnx_error}")
                                logger.info("📦 Using PyTorch JIT model for inference")
                        
                        if os.environ.get('AUTO_CONVERT_TO_ONNX_BACKGROUND', 'true').lower() == 'true':
                            threading.Thread(target=_convert_and_switch, daemon=True).start()
                        else:
                            _convert_and_switch()
                    
                    return True
                except Exception as jit_error:
                    logger.info(f"   Not a JIT model, trying as regular checkpoint... ({jit_error})")
                    try:
                        checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
                        self.is_jit_model = False
                    except Exception as load_error:
                        logger.error(f"Failed to load checkpoint: {load_error}")
                        return False
            else:
                # Load as regular checkpoint
                try:
                    checkpoint = torch.load(model_path, map_location='cpu', weights_only=False)
                    self.is_jit_model = False
                except Exception as load_error:
                    logger.error(f"Failed to load checkpoint: {load_error}")
                    logger.info("This may happen if PyTorch is not fully available in the .exe build")
                    return False
            
            # If we get here, it's not JIT, so load as checkpoint
            if not self.is_jit_model:
                try:
                    # Try to create the model - this might fail if nn.Module is None
                    self.model = FFCInpaintModel()
                    
                    if isinstance(checkpoint, dict):
                        if 'gen_state_dict' in checkpoint:
                            state_dict = checkpoint['gen_state_dict']
                            logger.info("📦 Found gen_state_dict")
                        elif 'state_dict' in checkpoint:
                            state_dict = checkpoint['state_dict']
                        elif 'model' in checkpoint:
                            state_dict = checkpoint['model']
                        else:
                            state_dict = checkpoint
                    else:
                        state_dict = checkpoint
                    
                    self._load_weights_with_mapping(self.model, state_dict)
                    
                    self.model.eval()
                    if self.use_gpu and self.device:
                        try:
                            self.model = self.model.to(self.device)
                        except Exception as gpu_error:
                            logger.warning(f"Could not move model to GPU: {gpu_error}")
                            logger.info("Using CPU instead")
                    
                    # Optional INT8 dynamic quantization for CPU eager model
                    if (self.int8_enabled or (self.quantize_enabled and not self.use_gpu and self.torch_precision in ('auto', 'int8'))) and not self.use_gpu:
                        try:
                            import torch.ao.quantization as tq  # type: ignore
                            self.model = tq.quantize_dynamic(self.model, {nn.Linear}, dtype=torch.qint8)  # type: ignore
                            logger.info("🔻 Applied dynamic INT8 quantization to inpainting model (CPU)")
                            self.torch_quantize_applied = True
                        except Exception as qe:
                            logger.warning(f"INT8 dynamic quantization not applied: {qe}")
                            
                except Exception as model_error:
                    logger.error(f"Failed to create or initialize model: {model_error}")
                    logger.info("This may happen if PyTorch neural network modules are not available in the .exe build")
                    return False
            
            self.model_loaded = True
            self.current_method = method
            
            self.config[f'{method}_model_path'] = model_path
            self._save_config()
            
            logger.info(f"✅ {method.upper()} loaded!")
            
            # ONNX CONVERSION (optionally in background)
            if AUTO_CONVERT_TO_ONNX and model_path.endswith('.pt') and self.model_loaded:
                def _convert_and_switch():
                    try:
                        onnx_path = self.convert_to_onnx(model_path, method)
                        if onnx_path and self.load_onnx_model(onnx_path):
                            logger.info("🚀 Using ONNX model for inference")
                    except Exception as onnx_error:
                        logger.warning(f"ONNX conversion failed: {onnx_error}")
                        logger.info("📦 Continuing with PyTorch model")
                
                if os.environ.get('AUTO_CONVERT_TO_ONNX_BACKGROUND', 'true').lower() == 'true':
                    threading.Thread(target=_convert_and_switch, daemon=True).start()
                else:
                    _convert_and_switch()

            return True
            
        except Exception as e:
            logger.error(f"❌ Failed to load model: {e}")
            logger.error(traceback.format_exc())
            logger.info("Note: If running from .exe, some ML libraries may not be included")
            logger.info("This is normal for lightweight builds - inpainting will be disabled")
            self.model_loaded = False
            return False

    def load_model_with_retry(self, method, model_path, force_reload=False, retries: int = 2, retry_delay: float = 0.5) -> bool:
        """Attempt to load a model with retries.

        Returns True if loaded; False if all attempts fail. On failure, the inpainter will safely no-op.

        """
        try:
            attempts = max(0, int(retries)) + 1
        except Exception:
            attempts = 1
        for attempt in range(attempts):
            try:
                ok = self.load_model(method, model_path, force_reload=force_reload)
                if ok:
                    return True
            except Exception as e:
                logger.warning(f"Load attempt {attempt+1} failed with exception: {e}")
            # brief delay before next try
            if attempt < attempts - 1:
                try:
                    time.sleep(max(0.0, float(retry_delay)))
                except Exception:
                    pass
        # If we reach here, loading failed. Leave model unloaded so inpaint() no-ops and returns original image.
        logger.warning("All load attempts failed; local inpainting will fall back to returning original images (no-op)")
        self.model_loaded = False
        # Keep current_method for logging/context if provided
        try:
            self.current_method = method
        except Exception:
            pass
        return False

    def unload(self):
        """Release all heavy resources held by this inpainter instance."""
        try:
            # Release ONNX session and metadata
            try:
                if self.onnx_session is not None:
                    self.onnx_session = None
            except Exception:
                pass
            for attr in ['onnx_input_names', 'onnx_output_names', 'current_onnx_path', 'onnx_fixed_size']:
                try:
                    if hasattr(self, attr):
                        setattr(self, attr, None)
                except Exception:
                    pass

            # Release PyTorch model
            try:
                if self.model is not None:
                    if TORCH_AVAILABLE and torch is not None:
                        try:
                            # Move to CPU then drop reference
                            self.model = self.model.to('cpu') if hasattr(self.model, 'to') else None
                        except Exception:
                            pass
                    self.model = None
            except Exception:
                pass

            # Drop bubble detector reference (not the global cache)
            try:
                self.bubble_detector = None
            except Exception:
                pass

            # Update flags
            self.model_loaded = False
            self.use_onnx = False
            self.is_jit_model = False

            # Free CUDA cache and trigger GC
            try:
                if TORCH_AVAILABLE and torch is not None and torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except Exception:
                pass
            try:
                import gc
                gc.collect()
            except Exception:
                pass
        except Exception:
            # Never raise from unload
            pass
    
    def pad_img_to_modulo(self, img: np.ndarray, mod: int) -> Tuple[np.ndarray, Tuple[int, int, int, int]]:
        """Pad image to be divisible by mod"""
        if len(img.shape) == 2:
            height, width = img.shape
        else:
            height, width = img.shape[:2]
        
        pad_height = (mod - height % mod) % mod
        pad_width = (mod - width % mod) % mod
        
        pad_top = pad_height // 2
        pad_bottom = pad_height - pad_top
        pad_left = pad_width // 2
        pad_right = pad_width - pad_left
        
        if len(img.shape) == 2:
            padded = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right)), mode='reflect')
        else:
            padded = np.pad(img, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), mode='reflect')
        
        return padded, (pad_top, pad_bottom, pad_left, pad_right)
    
    def remove_padding(self, img: np.ndarray, padding: Tuple[int, int, int, int]) -> np.ndarray:
        """Remove padding from image"""
        pad_top, pad_bottom, pad_left, pad_right = padding
        
        if len(img.shape) == 2:
            return img[pad_top:img.shape[0]-pad_bottom, pad_left:img.shape[1]-pad_right]
        else:
            return img[pad_top:img.shape[0]-pad_bottom, pad_left:img.shape[1]-pad_right, :]

    def _inpaint_tiled(self, image, mask, tile_size, overlap, refinement='normal'):
        """Process image in tiles"""
        orig_h, orig_w = image.shape[:2]
        result = image.copy()
        
        # Calculate tile positions
        for y in range(0, orig_h, tile_size - overlap):
            for x in range(0, orig_w, tile_size - overlap):
                # Calculate tile boundaries
                x_end = min(x + tile_size, orig_w)
                y_end = min(y + tile_size, orig_h)
                
                # Adjust start to ensure full tile size if possible
                if x_end - x < tile_size and x > 0:
                    x = max(0, x_end - tile_size)
                if y_end - y < tile_size and y > 0:
                    y = max(0, y_end - tile_size)
                
                # Extract tile
                tile_img = image[y:y_end, x:x_end]
                tile_mask = mask[y:y_end, x:x_end]
                
                # Skip if no inpainting needed
                if np.sum(tile_mask) == 0:
                    continue
                
                # Process this tile with the actual model
                processed_tile = self._process_single_tile(tile_img, tile_mask, tile_size, refinement)

                # Auto-retry for tile if no visible change
                try:
                    if self._is_noop(tile_img, processed_tile, tile_mask):
                        kernel = np.ones((3, 3), np.uint8)
                        expanded = cv2.dilate(tile_mask, kernel, iterations=1)
                        processed_retry = self._process_single_tile(tile_img, expanded, tile_size, 'fast')
                        if self._is_noop(tile_img, processed_retry, expanded):
                            logger.warning("Tile remained unchanged after retry; proceeding without further fallback")
                            processed_tile = processed_retry
                        else:
                            processed_tile = processed_retry
                except Exception as e:
                    logger.debug(f"Tiled no-op detection error: {e}")
                
                # Blend tile back into result
                if overlap > 0 and (x > 0 or y > 0):
                    result[y:y_end, x:x_end] = self._blend_tile(
                        result[y:y_end, x:x_end],
                        processed_tile,
                        x > 0,
                        y > 0,
                        overlap
                    )
                else:
                    result[y:y_end, x:x_end] = processed_tile
        
        logger.info(f"✅ Tiled inpainting complete ({orig_w}x{orig_h} in {tile_size}x{tile_size} tiles)")
        time.sleep(0.1)  # Brief pause for stability
        logger.debug("💤 Tiled inpainting completion pausing briefly for stability")
        return result

    def _process_single_tile(self, tile_img, tile_mask, tile_size, refinement):
        """Process a single tile without tiling"""
        # Temporarily disable tiling
        old_tiling = self.tiling_enabled
        self.tiling_enabled = False
        result = self.inpaint(tile_img, tile_mask, refinement, _skip_hd=True)
        self.tiling_enabled = old_tiling
        return result

    def _blend_tile(self, existing, new_tile, blend_x, blend_y, overlap):
        """Blend a tile with existing result"""
        if not blend_x and not blend_y:
            # No blending needed for first tile
            return new_tile
        
        h, w = new_tile.shape[:2]
        result = new_tile.copy()
        
        # Create blend weights
        if blend_x and overlap > 0 and w > overlap:
            # Horizontal blend on left edge
            for i in range(overlap):
                alpha = i / overlap
                result[:, i] = existing[:, i] * (1 - alpha) + new_tile[:, i] * alpha
        
        if blend_y and overlap > 0 and h > overlap:
            # Vertical blend on top edge
            for i in range(overlap):
                alpha = i / overlap
                result[i, :] = existing[i, :] * (1 - alpha) + new_tile[i, :] * alpha
        
        return result

    def _is_noop(self, original: np.ndarray, result: np.ndarray, mask: np.ndarray, threshold: float = 0.75) -> bool:
        """Return True if inpainting produced negligible change within the masked area."""
        try:
            if original is None or result is None:
                return True
            if original.shape != result.shape:
                return False
            # Normalize mask to single channel boolean
            if mask is None:
                return False
            if len(mask.shape) == 3:
                mask_gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
            else:
                mask_gray = mask
            m = mask_gray > 0
            if not np.any(m):
                return False
            # Fast path
            if np.array_equal(original, result):
                return True
            diff = cv2.absdiff(result, original)
            if len(diff.shape) == 3:
                diff_gray = cv2.cvtColor(diff, cv2.COLOR_BGR2GRAY)
            else:
                diff_gray = diff
            mean_diff = float(np.mean(diff_gray[m]))
            return mean_diff < threshold
        except Exception as e:
            logger.debug(f"No-op detection failed: {e}")
            return False

    def _is_white_paste(self, result: np.ndarray, mask: np.ndarray, white_threshold: int = 245, ratio: float = 0.90) -> bool:
        """Detect 'white paste' failure: masked area mostly saturated near white."""
        try:
            if result is None or mask is None:
                return False
            if len(mask.shape) == 3:
                mask_gray = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
            else:
                mask_gray = mask
            m = mask_gray > 0
            if not np.any(m):
                return False
            if len(result.shape) == 3:
                white = (result[..., 0] >= white_threshold) & (result[..., 1] >= white_threshold) & (result[..., 2] >= white_threshold)
            else:
                white = result >= white_threshold
            count_mask = int(np.count_nonzero(m))
            count_white = int(np.count_nonzero(white & m))
            if count_mask == 0:
                return False
            frac = count_white / float(count_mask)
            return frac >= ratio
        except Exception as e:
            logger.debug(f"White paste detection failed: {e}")
            return False

    def _log_inpaint_diag(self, path: str, result: np.ndarray, mask: np.ndarray):
        try:
            h, w = result.shape[:2]
            if len(result.shape) == 3:
                stats = (float(result.min()), float(result.max()), float(result.mean()))
            else:
                stats = (float(result.min()), float(result.max()), float(result.mean()))
            logger.info(f"[Diag] Path={path} onnx_quant={self.onnx_quantize_applied} torch_quant={self.torch_quantize_applied} size={w}x{h} stats(min,max,mean)={stats}")
            if self._is_white_paste(result, mask):
                logger.warning(f"[Diag] White-paste detected (mask>0 mostly white)")
        except Exception as e:
            logger.debug(f"Diag log failed: {e}")

    def inpaint(self, image, mask, refinement='normal', _retry_attempt: int = 0, _skip_hd: bool = False, _skip_tiling: bool = False):
        """Inpaint - compatible with JIT, checkpoint, and ONNX models

        Implements HD strategy (Resize/Crop) similar to comic-translate to speed up large images.

        """
        # Check for stop at start
        if self._check_stop():
            self._log("⏹️ Inpainting stopped by user", "warning")
            return image
        
        if not self.model_loaded:
            self._log("No model loaded", "error")
            return image
        
        try:
            # Store original dimensions
            orig_h, orig_w = image.shape[:2]
            
            # HD strategy (mirror of comic-translate): optional RESIZE or CROP before core inpainting
            if not _skip_hd:
                try:
                    strategy = getattr(self, 'hd_strategy', 'resize') or 'resize'
                except Exception:
                    strategy = 'resize'
                H, W = orig_h, orig_w
                if strategy == 'resize' and max(H, W) > max(16, int(getattr(self, 'hd_resize_limit', 1536))):
                    limit = max(16, int(getattr(self, 'hd_resize_limit', 1536)))
                    ratio = float(limit) / float(max(H, W))
                    new_w = max(1, int(W * ratio + 0.5))
                    new_h = max(1, int(H * ratio + 0.5))
                    image_small = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LANCZOS4)
                    mask_small = mask if len(mask.shape) == 2 else cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
                    mask_small = cv2.resize(mask_small, (new_w, new_h), interpolation=cv2.INTER_NEAREST)
                    result_small = self.inpaint(image_small, mask_small, refinement, 0, _skip_hd=True, _skip_tiling=True)
                    result_full = cv2.resize(result_small, (W, H), interpolation=cv2.INTER_LANCZOS4)
                    # Paste only masked area
                    mask_gray = mask_small  # already gray but at small size
                    mask_gray = cv2.resize(mask_gray, (W, H), interpolation=cv2.INTER_NEAREST)
                    m = mask_gray > 0
                    out = image.copy()
                    out[m] = result_full[m]
                    return out
                elif strategy == 'crop' and max(H, W) > max(16, int(getattr(self, 'hd_crop_trigger_size', 1024))):
                    mask_gray0 = mask if len(mask.shape) == 2 else cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
                    _, thresh = cv2.threshold(mask_gray0, 127, 255, cv2.THRESH_BINARY)
                    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                    if contours:
                        out = image.copy()
                        margin = max(0, int(getattr(self, 'hd_crop_margin', 16)))
                        for cnt in contours:
                            x, y, w, h = cv2.boundingRect(cnt)
                            l = max(0, x - margin); t = max(0, y - margin)
                            r = min(W, x + w + margin); b = min(H, y + h + margin)
                            if r <= l or b <= t:
                                continue
                            crop_img = image[t:b, l:r]
                            crop_mask = mask_gray0[t:b, l:r]
                            patch = self.inpaint(crop_img, crop_mask, refinement, 0, _skip_hd=True, _skip_tiling=True)
                            out[t:b, l:r] = patch
                        return out
            
            if len(mask.shape) == 3:
                mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
            
            # Apply dilation for anime method
            if self.current_method == 'anime':
                kernel = np.ones((7, 7), np.uint8)
                mask = cv2.dilate(mask, kernel, iterations=1)

            # Use instance tiling settings for ALL models
            logger.info(f"🔍 Tiling check: enabled={self.tiling_enabled}, tile_size={self.tile_size}, image_size={orig_h}x{orig_w}")

            # If tiling is enabled and image is larger than tile size
            if (not _skip_tiling) and self.tiling_enabled and (orig_h > self.tile_size or orig_w > self.tile_size):
                logger.info(f"🔲 Using tiled inpainting: {self.tile_size}x{self.tile_size} tiles with {self.tile_overlap}px overlap")
                return self._inpaint_tiled(image, mask, self.tile_size, self.tile_overlap, refinement)
                
            # ONNX inference path
            if self.use_onnx and self.onnx_session:
                logger.debug("Using ONNX inference")
                
                # CRITICAL: Convert BGR (OpenCV default) to RGB (ML model expected)
                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                
                # Check if this is a Carve model
                is_carve_model = False
                if hasattr(self, 'current_onnx_path'):
                    is_carve_model = "lama_fp32" in self.current_onnx_path or "carve" in self.current_onnx_path.lower()
                
                # Handle fixed-size models (resize instead of padding)
                if hasattr(self, 'onnx_fixed_size') and self.onnx_fixed_size:
                    fixed_h, fixed_w = self.onnx_fixed_size
                    # Resize to fixed size
                    image_resized = cv2.resize(image_rgb, (fixed_w, fixed_h), interpolation=cv2.INTER_LANCZOS4)
                    mask_resized = cv2.resize(mask, (fixed_w, fixed_h), interpolation=cv2.INTER_NEAREST)
                    
                    # Prepare inputs based on model type
                    if is_carve_model:
                        # Carve model expects normalized input [0, 1] range
                        logger.debug("Using Carve model normalization [0, 1]")
                        img_np = image_resized.astype(np.float32) / 255.0
                        mask_np = mask_resized.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0.5) * 1.0  # Binary mask
                    elif self.current_method == 'aot' or 'aot' in str(self.current_method).lower():
                        # AOT normalization: [-1, 1] range for image
                        logger.debug("Using AOT model normalization [-1, 1] for image, [0, 1] for mask")
                        img_np = (image_resized.astype(np.float32) / 127.5) - 1.0
                        mask_np = mask_resized.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0.5) * 1.0  # Binary mask
                        img_np = img_np * (1 - mask_np[:, :, np.newaxis])  # Mask out regions
                    elif 'anime' in str(self.current_method).lower():
                        # Anime/Manga LaMa normalization: [0, 1] range with optional input masking for stability
                        logger.debug("Using Anime/Manga LaMa normalization [0, 1] with input masking")
                        img_np = image_resized.astype(np.float32) / 255.0
                        mask_np = mask_resized.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0.5) * 1.0  # Binary mask
                        # CRITICAL: Mask out input regions for better text region stability
                        # This helps the model focus on generating content rather than being influenced by text artifacts
                        img_np = img_np * (1 - mask_np[:, :, np.newaxis])
                    else:
                        # Standard LaMa normalization: [0, 1] range
                        logger.debug("Using standard LaMa normalization [0, 1]")
                        img_np = image_resized.astype(np.float32) / 255.0
                        mask_np = mask_resized.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0) * 1.0
                    
                    # Convert to NCHW format
                    img_np = img_np.transpose(2, 0, 1)[np.newaxis, ...]
                    mask_np = mask_np[np.newaxis, np.newaxis, ...]
                    
                    # Run ONNX inference
                    ort_inputs = {
                        self.onnx_input_names[0]: img_np.astype(np.float32),
                        self.onnx_input_names[1]: mask_np.astype(np.float32)
                    }
                    
                    ort_outputs = self.onnx_session.run(self.onnx_output_names, ort_inputs)
                    output = ort_outputs[0]
                    
                    # Post-process output based on model type
                    if is_carve_model:
                        # CRITICAL: Carve model outputs values ALREADY in [0, 255] range!
                        # DO NOT multiply by 255 or apply any scaling
                        logger.debug("Carve model output is already in [0, 255] range")
                        raw_output = output[0].transpose(1, 2, 0)
                        logger.debug(f"Carve output stats: min={raw_output.min():.3f}, max={raw_output.max():.3f}, mean={raw_output.mean():.3f}")
                        result = raw_output  # Just transpose, no scaling
                    elif self.current_method == 'aot' or 'aot' in str(self.current_method).lower():
                        # AOT: [-1, 1] to [0, 255]
                        result = ((output[0].transpose(1, 2, 0) + 1.0) * 127.5)
                    else:
                        # Standard: [0, 1] to [0, 255]
                        result = output[0].transpose(1, 2, 0) * 255
                    
                    result = np.clip(np.round(result), 0, 255).astype(np.uint8)
                    # CRITICAL: Convert RGB (model output) back to BGR (OpenCV expected)
                    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
                    
                    # Resize back to original size
                    result = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
                    self._log_inpaint_diag('onnx-fixed', result, mask)
                    
                else:
                    # Variable-size models (use padding)
                    image_padded, padding = self.pad_img_to_modulo(image_rgb, self.pad_mod)
                    mask_padded, _ = self.pad_img_to_modulo(mask, self.pad_mod)
                    
                    # Prepare inputs based on model type
                    if is_carve_model:
                        # Carve model normalization [0, 1]
                        logger.debug("Using Carve model normalization [0, 1]")
                        img_np = image_padded.astype(np.float32) / 255.0
                        mask_np = mask_padded.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0.5) * 1.0
                    elif self.current_method == 'aot' or 'aot' in str(self.current_method).lower():
                        # AOT normalization: [-1, 1] for image
                        logger.debug("Using AOT model normalization [-1, 1] for image, [0, 1] for mask")
                        img_np = (image_padded.astype(np.float32) / 127.5) - 1.0
                        mask_np = mask_padded.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0.5) * 1.0
                        img_np = img_np * (1 - mask_np[:, :, np.newaxis])  # Mask out regions
                    elif 'anime' in str(self.current_method).lower():
                        # Anime/Manga LaMa normalization: [0, 1] range with optional input masking for stability
                        logger.debug("Using Anime/Manga LaMa normalization [0, 1] with input masking")
                        img_np = image_padded.astype(np.float32) / 255.0
                        mask_np = mask_padded.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0.5) * 1.0  # Binary mask
                        # CRITICAL: Mask out input regions for better text region stability
                        # This helps the model focus on generating content rather than being influenced by text artifacts
                        img_np = img_np * (1 - mask_np[:, :, np.newaxis])
                    else:
                        # Standard LaMa normalization: [0, 1]
                        logger.debug("Using standard LaMa normalization [0, 1]")
                        img_np = image_padded.astype(np.float32) / 255.0
                        mask_np = mask_padded.astype(np.float32) / 255.0
                        mask_np = (mask_np > 0) * 1.0
                    
                    # Convert to NCHW format
                    img_np = img_np.transpose(2, 0, 1)[np.newaxis, ...]
                    mask_np = mask_np[np.newaxis, np.newaxis, ...]
                    
                    # Check for stop before inference
                    if self._check_stop():
                        self._log("⏹️ ONNX inference stopped by user", "warning")
                        return image
                    
                    # Run ONNX inference
                    ort_inputs = {
                        self.onnx_input_names[0]: img_np.astype(np.float32),
                        self.onnx_input_names[1]: mask_np.astype(np.float32)
                    }
                    
                    ort_outputs = self.onnx_session.run(self.onnx_output_names, ort_inputs)
                    output = ort_outputs[0]
                    
                    # Post-process output
                    if is_carve_model:
                        # CRITICAL: Carve model outputs values ALREADY in [0, 255] range!
                        logger.debug("Carve model output is already in [0, 255] range")
                        raw_output = output[0].transpose(1, 2, 0)
                        logger.debug(f"Carve output stats: min={raw_output.min():.3f}, max={raw_output.max():.3f}, mean={raw_output.mean():.3f}")
                        result = raw_output  # Just transpose, no scaling
                    elif self.current_method == 'aot' or 'aot' in str(self.current_method).lower():
                        result = ((output[0].transpose(1, 2, 0) + 1.0) * 127.5)
                    else:
                        result = output[0].transpose(1, 2, 0) * 255
                    
                    result = np.clip(np.round(result), 0, 255).astype(np.uint8)
                    # CRITICAL: Convert RGB (model output) back to BGR (OpenCV expected)
                    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
                    
                    # Remove padding
                    result = self.remove_padding(result, padding)
                    self._log_inpaint_diag('onnx-padded', result, mask)
                    
            elif self.is_jit_model:
                # JIT model processing
                if self.current_method == 'aot':
                    # Special handling for AOT model
                    logger.debug("Using AOT-specific preprocessing")
                    
                    # CRITICAL: Convert BGR (OpenCV) to RGB (AOT model expected)
                    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    
                    # Pad images to be divisible by mod
                    image_padded, padding = self.pad_img_to_modulo(image_rgb, self.pad_mod)
                    mask_padded, _ = self.pad_img_to_modulo(mask, self.pad_mod)
                    
                    # AOT normalization: [-1, 1] range
                    img_torch = torch.from_numpy(image_padded).permute(2, 0, 1).unsqueeze_(0).float() / 127.5 - 1.0
                    mask_torch = torch.from_numpy(mask_padded).unsqueeze_(0).unsqueeze_(0).float() / 255.0
                    
                    # Binarize mask for AOT
                    mask_torch[mask_torch < 0.5] = 0
                    mask_torch[mask_torch >= 0.5] = 1
                    
                    # Move to device
                    img_torch = img_torch.to(self.device)
                    mask_torch = mask_torch.to(self.device)
                    
                    # Optional FP16 on GPU for lower VRAM
                    if self.quantize_enabled and self.use_gpu:
                        try:
                            if self.torch_precision == 'fp16' or self.torch_precision == 'auto':
                                img_torch = img_torch.half()
                                mask_torch = mask_torch.half()
                        except Exception:
                            pass
                    
                    # CRITICAL FOR AOT: Apply mask to input image
                    img_torch = img_torch * (1 - mask_torch)
                    
                    logger.debug(f"AOT Image shape: {img_torch.shape}, Mask shape: {mask_torch.shape}")
                    
                    # Run inference
                    with torch.no_grad():
                        inpainted = self.model(img_torch, mask_torch)
                    
                    # Post-process AOT output: denormalize from [-1, 1] to [0, 255]
                    result = ((inpainted.cpu().squeeze_(0).permute(1, 2, 0).numpy() + 1.0) * 127.5)
                    result = np.clip(np.round(result), 0, 255).astype(np.uint8)
                    
                    # CRITICAL: Convert RGB (model output) back to BGR (OpenCV expected)
                    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
                    
                    # Remove padding
                    result = self.remove_padding(result, padding)
                    self._log_inpaint_diag('jit-aot', result, mask)
                    
                else:
                    # LaMa/Anime model processing
                    logger.debug(f"Using standard processing for {self.current_method}")
                    
                    # CRITICAL: Convert BGR (OpenCV) to RGB (LaMa/JIT models expected)
                    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                    
                    # Pad images to be divisible by mod
                    image_padded, padding = self.pad_img_to_modulo(image_rgb, self.pad_mod)
                    mask_padded, _ = self.pad_img_to_modulo(mask, self.pad_mod)
                    
                    # CRITICAL: Normalize to [0, 1] range for LaMa models
                    image_norm = image_padded.astype(np.float32) / 255.0
                    mask_norm = mask_padded.astype(np.float32) / 255.0
                    
                    # Binary mask (values > 0 become 1)
                    mask_binary = (mask_norm > 0) * 1.0
                    
                    # For anime models: mask out input regions for better text stability
                    if 'anime' in str(self.current_method).lower():
                        logger.debug("Applying input masking for anime model (text region stability)")
                        image_norm = image_norm * (1 - mask_binary[:, :, np.newaxis])
                    
                    # Convert to PyTorch tensors with correct shape
                    # Image should be [B, C, H, W]  
                    image_tensor = torch.from_numpy(image_norm).permute(2, 0, 1).unsqueeze(0).float()
                    mask_tensor = torch.from_numpy(mask_binary).unsqueeze(0).unsqueeze(0).float()
                    
                    # Move to device
                    image_tensor = image_tensor.to(self.device)
                    mask_tensor = mask_tensor.to(self.device)
                    
                    # Optional FP16 on GPU for lower VRAM
                    if self.quantize_enabled and self.use_gpu:
                        try:
                            if self.torch_precision == 'fp16' or self.torch_precision == 'auto':
                                image_tensor = image_tensor.half()
                                mask_tensor = mask_tensor.half()
                        except Exception:
                            pass
                    
                    # Debug shapes
                    logger.debug(f"Image tensor shape: {image_tensor.shape}")  # Should be [1, 3, H, W]
                    logger.debug(f"Mask tensor shape: {mask_tensor.shape}")    # Should be [1, 1, H, W]
                    
                    # Ensure spatial dimensions match
                    if image_tensor.shape[2:] != mask_tensor.shape[2:]:
                        logger.warning(f"Spatial dimension mismatch: image {image_tensor.shape[2:]}, mask {mask_tensor.shape[2:]}")
                        # Resize mask to match image
                        mask_tensor = F.interpolate(mask_tensor, size=image_tensor.shape[2:], mode='nearest')
                    
                    # Run inference with proper error handling
                    with torch.no_grad():
                        try:
                            # Standard LaMa JIT models expect (image, mask)
                            inpainted = self.model(image_tensor, mask_tensor)
                        except RuntimeError as e:
                            error_str = str(e)
                            logger.error(f"Model inference failed: {error_str}")
                            
                            # If tensor size mismatch, log detailed info
                            if "size of tensor" in error_str.lower():
                                logger.error(f"Image shape: {image_tensor.shape}")
                                logger.error(f"Mask shape: {mask_tensor.shape}")
                                
                                # Try transposing if needed
                                if "dimension 3" in error_str and "880" in error_str:
                                    # This suggests the tensors might be in wrong format
                                    # Try different permutation
                                    logger.info("Attempting to fix tensor format...")
                                    
                                    # Ensure image is [B, C, H, W] not [B, H, W, C]
                                    if image_tensor.shape[1] > 3:
                                        image_tensor = image_tensor.permute(0, 3, 1, 2)
                                        logger.info(f"Permuted image to: {image_tensor.shape}")
                                    
                                    # Try again
                                    inpainted = self.model(image_tensor, mask_tensor)
                                else:
                                    # As last resort, try swapped arguments
                                    logger.info("Trying swapped arguments (mask, image)...")
                                    inpainted = self.model(mask_tensor, image_tensor)
                            else:
                                raise e
                    
                    # Process output
                    # Output should be [B, C, H, W]
                    if len(inpainted.shape) == 4:
                        # Remove batch dimension and permute to [H, W, C]
                        result = inpainted[0].permute(1, 2, 0).detach().cpu().numpy()
                    else:
                        # Handle unexpected output shape
                        result = inpainted.detach().cpu().numpy()
                        if len(result.shape) == 3 and result.shape[0] == 3:
                            result = result.transpose(1, 2, 0)
                    
                    # Denormalize to 0-255 range
                    result = np.clip(result * 255, 0, 255).astype(np.uint8)
                    
                    # CRITICAL: Convert RGB (model output) back to BGR (OpenCV expected)
                    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
                    
                    # Remove padding
                    result = self.remove_padding(result, padding)
                    self._log_inpaint_diag('jit-lama', result, mask)
            
            else:
                # Original checkpoint model processing (keep as is)
                h, w = image.shape[:2]
                size = 768 if self.current_method == 'anime' else 512
                
                img_resized = cv2.resize(image, (size, size), interpolation=cv2.INTER_LANCZOS4)
                mask_resized = cv2.resize(mask, (size, size), interpolation=cv2.INTER_NEAREST)
                
                img_norm = img_resized.astype(np.float32) / 127.5 - 1
                mask_norm = mask_resized.astype(np.float32) / 255.0
                
                img_tensor = torch.from_numpy(img_norm).permute(2, 0, 1).unsqueeze(0).float()
                mask_tensor = torch.from_numpy(mask_norm).unsqueeze(0).unsqueeze(0).float()
                
                if self.use_gpu and self.device:
                    img_tensor = img_tensor.to(self.device)
                    mask_tensor = mask_tensor.to(self.device)
                
                with torch.no_grad():
                    output = self.model(img_tensor, mask_tensor)
                
                result = output.squeeze(0).permute(1, 2, 0).cpu().numpy()
                result = ((result + 1) * 127.5).clip(0, 255).astype(np.uint8)
                result = cv2.resize(result, (w, h), interpolation=cv2.INTER_LANCZOS4)
                self._log_inpaint_diag('ckpt', result, mask)
            
            # Ensure result matches original size exactly
            if result.shape[:2] != (orig_h, orig_w):
                result = cv2.resize(result, (orig_w, orig_h), interpolation=cv2.INTER_LANCZOS4)
            
            # Apply refinement blending if requested
            if refinement != 'fast':
                # Ensure mask is same size as result
                if mask.shape[:2] != (orig_h, orig_w):
                    mask = cv2.resize(mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
                
                mask_3ch = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) / 255.0
                kernel = cv2.getGaussianKernel(21, 5)
                kernel = kernel @ kernel.T
                mask_blur = cv2.filter2D(mask_3ch, -1, kernel)
                result = (result * mask_blur + image * (1 - mask_blur)).astype(np.uint8)
            
            # No-op detection and auto-retry
            try:
                if self._is_noop(image, result, mask):
                    if _retry_attempt == 0:
                        logger.warning("⚠️ Inpainting produced no visible change; retrying with slight mask dilation and fast refinement")
                        kernel = np.ones((3, 3), np.uint8)
                        expanded_mask = cv2.dilate(mask, kernel, iterations=1)
                        return self.inpaint(image, expanded_mask, refinement='fast', _retry_attempt=1)
                    elif _retry_attempt == 1:
                        logger.warning("⚠️ Still no visible change after retry; attempting a second dilation and fast refinement")
                        kernel = np.ones((5, 5), np.uint8)
                        expanded_mask2 = cv2.dilate(mask, kernel, iterations=1)
                        return self.inpaint(image, expanded_mask2, refinement='fast', _retry_attempt=2)
                    else:
                        logger.warning("⚠️ No further retries; returning last result without fallback")
            except Exception as e:
                logger.debug(f"No-op detection step failed: {e}")
            
            logger.info("✅ Inpainted successfully!")
            
            # Force garbage collection to reduce memory spikes
            try:
                import gc
                gc.collect()
                # Clear CUDA cache if using GPU
                if torch is not None and hasattr(torch, 'cuda') and torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except Exception:
                pass
            
            time.sleep(0.1)  # Brief pause for stability
            logger.debug("💤 Inpainting completion pausing briefly for stability")
            return result
            
        except Exception as e:
            logger.error(f"❌ Inpainting failed: {e}")
            logger.error(traceback.format_exc())
            
            # Return original image on failure
            logger.warning("Returning original image due to error")
            return image
    
    def inpaint_with_prompt(self, image, mask, prompt=None):
        """Compatibility method"""
        return self.inpaint(image, mask)
    
    def batch_inpaint(self, images, masks):
        """Batch inpainting"""
        return [self.inpaint(img, mask) for img, mask in zip(images, masks)]
    
    def load_bubble_model(self, model_path: str) -> bool:
        """Load bubble detection model"""
        if not BUBBLE_DETECTOR_AVAILABLE:
            logger.warning("Bubble detector not available")
            return False
        
        if self.bubble_detector is None:
            self.bubble_detector = BubbleDetector()
        
        if self.bubble_detector.load_model(model_path):
            self.bubble_model_loaded = True
            self.config['bubble_model_path'] = model_path
            self._save_config()
            logger.info("✅ Bubble detection model loaded")
            return True
        
        return False
    
    def detect_bubbles(self, image_path: str, confidence: float = 0.5) -> List[Tuple[int, int, int, int]]:
        """Detect speech bubbles in image"""
        if not self.bubble_model_loaded or self.bubble_detector is None:
            logger.warning("No bubble model loaded")
            return []
        
        return self.bubble_detector.detect_bubbles(image_path, confidence=confidence)
    
    def create_bubble_mask(self, image: np.ndarray, bubbles: List[Tuple[int, int, int, int]], 

                          expand_pixels: int = 5) -> np.ndarray:
        """Create mask from detected bubbles"""
        h, w = image.shape[:2]
        mask = np.zeros((h, w), dtype=np.uint8)
        
        for x, y, bw, bh in bubbles:
            x1 = max(0, x - expand_pixels)
            y1 = max(0, y - expand_pixels)
            x2 = min(w, x + bw + expand_pixels)
            y2 = min(h, y + bh + expand_pixels)
            
            cv2.rectangle(mask, (x1, y1), (x2, y2), 255, -1)
        
        return mask
    
    def inpaint_with_bubble_detection(self, image_path: str, confidence: float = 0.5,

                                     expand_pixels: int = 5, refinement: str = 'normal') -> np.ndarray:
        """Inpaint using automatic bubble detection"""
        image = cv2.imread(image_path)
        if image is None:
            logger.error(f"Failed to load image: {image_path}")
            return None
        
        bubbles = self.detect_bubbles(image_path, confidence)
        if not bubbles:
            logger.warning("No bubbles detected")
            return image
        
        logger.info(f"Detected {len(bubbles)} bubbles")
        
        mask = self.create_bubble_mask(image, bubbles, expand_pixels)
        result = self.inpaint(image, mask, refinement)
        
        return result
    
    def batch_inpaint_with_bubbles(self, image_paths: List[str], **kwargs) -> List[np.ndarray]:
        """Batch inpaint multiple images with bubble detection"""
        results = []
        
        for i, image_path in enumerate(image_paths):
            logger.info(f"Processing image {i+1}/{len(image_paths)}")
            result = self.inpaint_with_bubble_detection(image_path, **kwargs)
            results.append(result)
        
        return results


# Compatibility classes - MAINTAIN ALL ORIGINAL CLASSES
class LaMaModel(FFCInpaintModel):
    pass

class MATModel(FFCInpaintModel):
    pass

class AOTModel(FFCInpaintModel):
    pass

class SDInpaintModel(FFCInpaintModel):
    pass

class AnimeMangaInpaintModel(FFCInpaintModel):
    pass

class LaMaOfficialModel(FFCInpaintModel):
    pass


class HybridInpainter:
    """Hybrid inpainter for compatibility"""
    
    def __init__(self):
        self.inpainters = {}
    
    def add_method(self, name, method, model_path):
        """Add a method - maintains compatibility"""
        try:
            inpainter = LocalInpainter()
            if inpainter.load_model(method, model_path):
                self.inpainters[name] = inpainter
                return True
        except:
            pass
        return False
    
    def inpaint_ensemble(self, image: np.ndarray, mask: np.ndarray, 

                        weights: Dict[str, float] = None) -> np.ndarray:
        """Ensemble inpainting"""
        if not self.inpainters:
            logger.error("No inpainters loaded")
            return image
        
        if weights is None:
            weights = {name: 1.0 / len(self.inpainters) for name in self.inpainters}
        
        results = []
        for name, inpainter in self.inpainters.items():
            result = inpainter.inpaint(image, mask)
            weight = weights.get(name, 1.0 / len(self.inpainters))
            results.append(result * weight)
        
        ensemble = np.sum(results, axis=0).astype(np.uint8)
        return ensemble


# Helper function for quick setup
def setup_inpainter_for_manga(auto_download=True):
    """Quick setup for manga inpainting"""
    inpainter = LocalInpainter()
    
    if auto_download:
        # Try to download anime JIT model
        jit_path = inpainter.download_jit_model('anime')
        if jit_path:
            inpainter.load_model('anime', jit_path)
            logger.info("✅ Manga inpainter ready with JIT model")
    
    return inpainter


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1:
        if sys.argv[1] == "download_jit":
            # Download JIT models
            inpainter = LocalInpainter()
            for method in ['lama', 'anime', 'lama_official']:
                print(f"\nDownloading {method}...")
                path = inpainter.download_jit_model(method)
                if path:
                    print(f"  ✅ Downloaded to: {path}")
        
        elif len(sys.argv) > 2:
            # Test with model
            inpainter = LocalInpainter()
            inpainter.load_model('lama', sys.argv[1])
            print("Model loaded - check logs for details")
    
    else:
        print("\nLocal Inpainter - Compatible Version")
        print("=====================================")
        print("\nSupports both:")
        print("  - JIT models (.pt) - RECOMMENDED")
        print("  - Checkpoint files (.ckpt) - With warnings")
        print("\nTo download JIT models:")
        print("  python local_inpainter.py download_jit")
        print("\nTo test:")
        print("  python local_inpainter.py <model_path>")