Method load

modules/icon_picture_processor.py:63–121

Load RMBG-2.0 ONNX model; fallback to CPU if CUDA fails.

(self)

        )
    
    def load(self):
        """Load RMBG-2.0 ONNX model; fallback to CPU if CUDA fails."""
        if self._is_loaded:
            return
        
        if not ONNX_AVAILABLE:
            print("[RMBGModel] Warning: onnxruntime not available, using fallback mode")
            self._is_loaded = True
            return
        
        if not os.path.exists(self.model_path):
            print(f"[RMBGModel] Warning: Model file not found at {self.model_path}, using fallback mode")
            self._is_loaded = True
            return
        
        # ONNX Runtime options
        session_options = ort.SessionOptions()
        session_options.log_severity_level = 3  # ERROR only
        session_options.enable_profiling = False
        
        # Available providers
        available_providers = ort.get_available_providers()
        
        # Try CUDA then CPU
        providers_to_try = [
            (['CUDAExecutionProvider', 'CPUExecutionProvider'], "CUDA+CPU"),
            (['CPUExecutionProvider'], "CPU only"),
        ]
        
        for providers, name in providers_to_try:
            # Filter valid providers
            valid_providers = [p for p in providers if p in available_providers]
            if not valid_providers:
                continue
            
            try:
                print(f"[RMBGModel] Trying to load with {name} ({valid_providers})...")
                self._session = ort.InferenceSession(
                    self.model_path,
                    providers=valid_providers,
                    sess_options=session_options
                )
                
                self._input_name = self._session.get_inputs()[0].name
                self._output_name = self._session.get_outputs()[0].name
                self._providers = valid_providers
                
                self._is_loaded = True
                print(f"[RMBGModel] Model loaded successfully with {name}")
                return
                
            except Exception as e:
                print(f"[RMBGModel] Failed to load with {name}: {e}")
                # Try next config
                continue
        
        # All attempts failed, use fallback
        print("[RMBGModel] Warning: All loading attempts failed, using fallback mode (no background removal)")