629 lines
26 KiB
Python
629 lines
26 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
BMO MiniMax-Remover ComfyUI Node
|
||
High-quality video object removal - based on official implementation
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import torch
|
||
import numpy as np
|
||
import cv2
|
||
from typing import Optional, Union, List
|
||
import folder_paths
|
||
import comfy.model_management as model_management
|
||
|
||
# Import the BMO pipeline
|
||
from pathlib import Path
|
||
|
||
# Add current directory to path for imports
|
||
current_dir = os.path.dirname(os.path.abspath(__file__))
|
||
if current_dir not in sys.path:
|
||
sys.path.append(current_dir)
|
||
|
||
# Handle import path for ComfyUI
|
||
def get_comfyui_base_path():
|
||
"""Get the ComfyUI base path for proper imports"""
|
||
current_dir = Path(__file__).parent.absolute()
|
||
|
||
# Look for ComfyUI base directory
|
||
comfyui_paths = [
|
||
current_dir.parent.parent, # custom_nodes/minimax-remover-bmo/
|
||
Path("D:/COMFY_UI/ComfyUI"), # Direct path
|
||
Path(os.environ.get("COMFYUI_BASE", "")) # Environment variable
|
||
]
|
||
|
||
for path in comfyui_paths:
|
||
if path.exists() and (path / "nodes.py").exists():
|
||
return str(path)
|
||
|
||
return str(current_dir.parent.parent)
|
||
|
||
# Add ComfyUI to path
|
||
comfyui_base = get_comfyui_base_path()
|
||
if comfyui_base not in sys.path:
|
||
sys.path.insert(0, comfyui_base)
|
||
|
||
try:
|
||
# ComfyUI imports
|
||
import comfy.model_management as model_management
|
||
print("✅ ComfyUI model_management imported successfully")
|
||
COMFYUI_AVAILABLE = True
|
||
except ImportError:
|
||
print(f"⚠️ Could not import ComfyUI model_management")
|
||
COMFYUI_AVAILABLE = False
|
||
# Fallback device management
|
||
class MockModelManagement:
|
||
@staticmethod
|
||
def get_torch_device():
|
||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
model_management = MockModelManagement()
|
||
|
||
# Import required modules
|
||
from diffusers.models import AutoencoderKLWan
|
||
from diffusers.schedulers import UniPCMultistepScheduler
|
||
|
||
# Enhanced imports with path handling
|
||
current_dir = Path(__file__).parent
|
||
try:
|
||
# Try local import first
|
||
sys.path.insert(0, str(current_dir))
|
||
from pipeline_minimax_remover_bmo import Minimax_Remover_Pipeline_BMO
|
||
from transformer_minimax_remover import Transformer3DModel
|
||
print("✅ Local BMO modules imported successfully")
|
||
except ImportError as e:
|
||
print(f"❌ Failed to import BMO modules: {e}")
|
||
# Attempt to find and import from different locations
|
||
possible_paths = [
|
||
current_dir,
|
||
current_dir.parent,
|
||
Path(comfyui_base) / "custom_nodes" / "minimax-remover-bmo"
|
||
]
|
||
|
||
imported = False
|
||
for path in possible_paths:
|
||
try:
|
||
if str(path) not in sys.path:
|
||
sys.path.insert(0, str(path))
|
||
from pipeline_minimax_remover_bmo import Minimax_Remover_Pipeline_BMO
|
||
from transformer_minimax_remover import Transformer3DModel
|
||
print(f"✅ BMO modules imported from: {path}")
|
||
imported = True
|
||
break
|
||
except ImportError:
|
||
continue
|
||
|
||
if not imported:
|
||
raise ImportError("Could not import BMO pipeline modules from any location")
|
||
|
||
# Lazy imports for diffusers components
|
||
def lazy_import_diffusers():
|
||
"""Lazy import diffusers components to avoid import issues"""
|
||
global AutoencoderKLWan, UniPCMultistepScheduler, Minimax_Remover_Pipeline_BMO, Transformer3DModel
|
||
|
||
if 'AutoencoderKLWan' not in globals():
|
||
from diffusers.models import AutoencoderKLWan
|
||
from diffusers.schedulers import UniPCMultistepScheduler
|
||
from pipeline_minimax_remover_bmo import Minimax_Remover_Pipeline_BMO
|
||
from transformer_minimax_remover import Transformer3DModel
|
||
|
||
return AutoencoderKLWan, UniPCMultistepScheduler, Minimax_Remover_Pipeline_BMO, Transformer3DModel
|
||
|
||
class MinimaxRemoverBMONode:
|
||
"""
|
||
BMO MiniMax-Remover Node for ComfyUI
|
||
High-quality video object removal with separate model path inputs
|
||
"""
|
||
|
||
def __init__(self):
|
||
self.pipe = None
|
||
self.device = model_management.get_torch_device()
|
||
self.comfyui_models_path = Path(comfyui_base) / "models"
|
||
|
||
def auto_download_models(self, force_download=False):
|
||
"""
|
||
Automatically download models if they don't exist
|
||
Returns the paths to the downloaded models
|
||
"""
|
||
print("🔍 Checking for MiniMax-Remover models...")
|
||
|
||
# Standard model locations to try
|
||
possible_locations = [
|
||
# Project directory
|
||
Path("models"),
|
||
# ComfyUI models directory
|
||
self.comfyui_models_path,
|
||
# User's cache directory
|
||
Path.home() / ".cache" / "minimax-remover"
|
||
]
|
||
|
||
# Check if models already exist
|
||
for base_path in possible_locations:
|
||
vae_path = base_path / "vae"
|
||
transformer_path = base_path / "transformer"
|
||
scheduler_path = base_path / "scheduler"
|
||
|
||
if (vae_path.exists() and transformer_path.exists() and scheduler_path.exists() and
|
||
(vae_path / "config.json").exists() and
|
||
(transformer_path / "config.json").exists() and
|
||
(scheduler_path / "scheduler_config.json").exists()):
|
||
|
||
print(f"✅ Found existing models at: {base_path}")
|
||
return str(vae_path), str(transformer_path), str(scheduler_path)
|
||
|
||
if not force_download:
|
||
print("📥 Models not found locally. Starting automatic download...")
|
||
|
||
# Choose download location (prefer project directory, fallback to cache)
|
||
download_base = Path("models")
|
||
if not download_base.parent.exists() or not os.access(download_base.parent, os.W_OK):
|
||
download_base = Path.home() / ".cache" / "minimax-remover" / "models"
|
||
print(f"📁 Using cache directory for models: {download_base}")
|
||
|
||
download_base.mkdir(parents=True, exist_ok=True)
|
||
|
||
try:
|
||
print(f"🌐 Downloading MiniMax-Remover models to: {download_base}")
|
||
print(f"📊 Expected download size: ~25-30 GB")
|
||
print(f"⏳ This may take several minutes depending on your connection...")
|
||
|
||
# Import huggingface_hub for downloading
|
||
try:
|
||
from huggingface_hub import snapshot_download
|
||
except ImportError:
|
||
print("❌ huggingface_hub not found. Installing...")
|
||
import subprocess
|
||
subprocess.check_call([sys.executable, "-m", "pip", "install", "huggingface_hub"])
|
||
from huggingface_hub import snapshot_download
|
||
|
||
# Download with progress
|
||
snapshot_download(
|
||
repo_id="zibojia/minimax-remover",
|
||
local_dir=str(download_base),
|
||
local_dir_use_symlinks=False,
|
||
allow_patterns=["vae/*", "transformer/*", "scheduler/*"],
|
||
ignore_patterns=["*.git*", "README.md", "*.txt"]
|
||
)
|
||
|
||
# Verify download
|
||
vae_path = download_base / "vae"
|
||
transformer_path = download_base / "transformer"
|
||
scheduler_path = download_base / "scheduler"
|
||
|
||
if (vae_path.exists() and transformer_path.exists() and scheduler_path.exists()):
|
||
print("🎉 Models downloaded successfully!")
|
||
print(f"📁 VAE: {vae_path}")
|
||
print(f"📁 Transformer: {transformer_path}")
|
||
print(f"📁 Scheduler: {scheduler_path}")
|
||
return str(vae_path), str(transformer_path), str(scheduler_path)
|
||
else:
|
||
raise Exception("Download completed but models not found in expected locations")
|
||
|
||
except Exception as e:
|
||
print(f"❌ Auto-download failed: {e}")
|
||
print("\n🔧 Manual download options:")
|
||
print("1. Run: python download_models.py")
|
||
print("2. Run: huggingface-cli download zibojia/minimax-remover --local-dir ./models")
|
||
print("3. See MODEL_DOWNLOAD_GUIDE.md for detailed instructions")
|
||
|
||
# Return default paths so user can manually configure
|
||
return "models/vae/", "models/transformer/", "models/scheduler/"
|
||
|
||
@classmethod
|
||
def INPUT_TYPES(cls):
|
||
return {
|
||
"required": {
|
||
"images": ("IMAGE",),
|
||
"masks": ("MASK",),
|
||
"num_inference_steps": ("INT", {
|
||
"default": 12,
|
||
"min": 6,
|
||
"max": 50,
|
||
"step": 1,
|
||
"display": "number"
|
||
}),
|
||
"iterations": ("INT", {
|
||
"default": 6,
|
||
"min": 1,
|
||
"max": 20,
|
||
"step": 1,
|
||
"display": "number"
|
||
}),
|
||
"seed": ("INT", {
|
||
"default": 42,
|
||
"min": 0,
|
||
"max": 0xffffffffffffffff,
|
||
"display": "number"
|
||
}),
|
||
},
|
||
"optional": {
|
||
"auto_download": ("BOOLEAN", {
|
||
"default": True,
|
||
"tooltip": "Automatically download models if not found"
|
||
}),
|
||
"vae_path": ("STRING", {
|
||
"default": "auto",
|
||
"multiline": False,
|
||
"tooltip": "Path to VAE model directory (auto = automatic detection)"
|
||
}),
|
||
"transformer_path": ("STRING", {
|
||
"default": "auto",
|
||
"multiline": False,
|
||
"tooltip": "Path to Transformer model directory (auto = automatic detection)"
|
||
}),
|
||
"scheduler_path": ("STRING", {
|
||
"default": "auto",
|
||
"multiline": False,
|
||
"tooltip": "Path to Scheduler config directory (auto = automatic detection)"
|
||
}),
|
||
}
|
||
}
|
||
|
||
RETURN_TYPES = ("IMAGE",)
|
||
RETURN_NAMES = ("images",)
|
||
FUNCTION = "process_video"
|
||
CATEGORY = "MiniMax-Remover"
|
||
DESCRIPTION = "BMO MiniMax video object removal with separate model paths"
|
||
|
||
def resolve_model_path(self, model_path: str, model_type: str) -> str:
|
||
"""
|
||
Resolve model path with intelligent fallbacks
|
||
|
||
Args:
|
||
model_path: User-provided path (can be relative or absolute)
|
||
model_type: Type of model (vae, transformer, scheduler)
|
||
|
||
Returns:
|
||
Absolute path to model directory
|
||
"""
|
||
# Convert to Path object
|
||
path = Path(model_path)
|
||
|
||
# If absolute path and exists, use it
|
||
if path.is_absolute() and path.exists():
|
||
print(f"✅ Using absolute path for {model_type}: {path}")
|
||
return str(path)
|
||
|
||
# Try relative to ComfyUI base
|
||
comfyui_relative = Path(comfyui_base) / model_path
|
||
if comfyui_relative.exists():
|
||
print(f"✅ Using ComfyUI relative path for {model_type}: {comfyui_relative}")
|
||
return str(comfyui_relative)
|
||
|
||
# Try in ComfyUI models directory
|
||
models_path = self.comfyui_models_path / model_type
|
||
if models_path.exists():
|
||
print(f"✅ Using ComfyUI models path for {model_type}: {models_path}")
|
||
return str(models_path)
|
||
|
||
# Try original path as-is (fallback)
|
||
print(f"⚠️ Using fallback path for {model_type}: {path}")
|
||
return str(path)
|
||
|
||
def load_models(self, vae_path: str, transformer_path: str, scheduler_path: str, auto_download: bool = True):
|
||
"""Load MiniMax-Remover models from paths with auto-download support"""
|
||
if self.pipe is not None:
|
||
print("ℹ️ Models already loaded, skipping...")
|
||
return # Already loaded
|
||
|
||
print("🔄 Loading BMO MiniMax-Remover models...")
|
||
print(f"🗂️ ComfyUI base path: {comfyui_base}")
|
||
print(f"🗂️ ComfyUI models path: {self.comfyui_models_path}")
|
||
|
||
# Handle auto-download and path resolution
|
||
if auto_download and (vae_path == "auto" or transformer_path == "auto" or scheduler_path == "auto"):
|
||
print("🤖 Auto-download mode enabled")
|
||
try:
|
||
auto_vae, auto_transformer, auto_scheduler = self.auto_download_models()
|
||
|
||
# Use auto-detected paths for "auto" values
|
||
if vae_path == "auto":
|
||
vae_path = auto_vae
|
||
if transformer_path == "auto":
|
||
transformer_path = auto_transformer
|
||
if scheduler_path == "auto":
|
||
scheduler_path = auto_scheduler
|
||
|
||
print(f"🎯 Auto-resolved paths:")
|
||
print(f" VAE: {vae_path}")
|
||
print(f" Transformer: {transformer_path}")
|
||
print(f" Scheduler: {scheduler_path}")
|
||
|
||
except Exception as e:
|
||
print(f"⚠️ Auto-download failed, falling back to manual paths: {e}")
|
||
# Keep original paths if auto-download fails
|
||
|
||
try:
|
||
# Resolve model paths
|
||
resolved_vae_path = self.resolve_model_path(vae_path, "vae")
|
||
resolved_transformer_path = self.resolve_model_path(transformer_path, "transformer")
|
||
resolved_scheduler_path = self.resolve_model_path(scheduler_path, "scheduler")
|
||
|
||
print(f"📁 Loading VAE from: {resolved_vae_path}")
|
||
print(f"📁 Loading Transformer from: {resolved_transformer_path}")
|
||
print(f"📁 Loading Scheduler from: {resolved_scheduler_path}")
|
||
|
||
# Check if paths exist before attempting to load
|
||
missing_paths = []
|
||
for name, path in [("VAE", resolved_vae_path), ("Transformer", resolved_transformer_path), ("Scheduler", resolved_scheduler_path)]:
|
||
if not Path(path).exists():
|
||
missing_paths.append(f"{name}: {path}")
|
||
|
||
if missing_paths:
|
||
print(f"❌ Missing model paths:")
|
||
for missing in missing_paths:
|
||
print(f" {missing}")
|
||
|
||
if auto_download:
|
||
print("🔄 Attempting to re-download missing models...")
|
||
auto_vae, auto_transformer, auto_scheduler = self.auto_download_models(force_download=True)
|
||
resolved_vae_path = auto_vae
|
||
resolved_transformer_path = auto_transformer
|
||
resolved_scheduler_path = auto_scheduler
|
||
else:
|
||
raise FileNotFoundError(f"Models not found. Enable auto_download or check paths.")
|
||
|
||
# Load models from resolved paths
|
||
AutoencoderKLWan, UniPCMultistepScheduler, Minimax_Remover_Pipeline_BMO, Transformer3DModel = lazy_import_diffusers()
|
||
|
||
vae = AutoencoderKLWan.from_pretrained(
|
||
resolved_vae_path,
|
||
torch_dtype=torch.float16
|
||
)
|
||
transformer = Transformer3DModel.from_pretrained(
|
||
resolved_transformer_path,
|
||
torch_dtype=torch.float16
|
||
)
|
||
scheduler = UniPCMultistepScheduler.from_pretrained(
|
||
resolved_scheduler_path
|
||
)
|
||
|
||
# Create the BMO pipeline
|
||
self.pipe = Minimax_Remover_Pipeline_BMO(
|
||
vae=vae,
|
||
transformer=transformer,
|
||
scheduler=scheduler
|
||
).to(self.device)
|
||
|
||
print("✅ BMO MiniMax-Remover models loaded successfully!")
|
||
print(f" Using device: {self.device}")
|
||
print(f" VAE: {type(vae).__name__}")
|
||
print(f" Transformer: {type(transformer).__name__}")
|
||
print(f" Scheduler: {type(scheduler).__name__}")
|
||
|
||
except Exception as e:
|
||
print(f"❌ Failed to load models: {e}")
|
||
print("🔍 Debug info:")
|
||
print(f" VAE path exists: {Path(resolved_vae_path).exists()}")
|
||
print(f" Transformer path exists: {Path(resolved_transformer_path).exists()}")
|
||
print(f" Scheduler path exists: {Path(resolved_scheduler_path).exists()}")
|
||
|
||
# List available files for debugging
|
||
for name, path in [("VAE", resolved_vae_path), ("Transformer", resolved_transformer_path), ("Scheduler", resolved_scheduler_path)]:
|
||
if Path(path).exists():
|
||
files = list(Path(path).glob("*"))
|
||
print(f" {name} directory contents: {[f.name for f in files]}")
|
||
|
||
raise e
|
||
|
||
def process_video(
|
||
self,
|
||
images,
|
||
masks,
|
||
num_inference_steps=12,
|
||
iterations=6,
|
||
seed=42,
|
||
auto_download=True,
|
||
vae_path="auto",
|
||
transformer_path="auto",
|
||
scheduler_path="auto"
|
||
):
|
||
"""
|
||
Process video with BMO MiniMax-Remover with auto-download support
|
||
|
||
Args:
|
||
images: Input video frames [B, H, W, C] in [0, 1]
|
||
masks: Binary masks [B, H, W] in [0, 1]
|
||
num_inference_steps: Number of denoising steps (official default: 12)
|
||
iterations: Mask expansion iterations (official default: 6)
|
||
seed: Random seed for reproducible results
|
||
auto_download: Automatically download models if not found
|
||
vae_path: Path to VAE model directory (auto = automatic detection)
|
||
transformer_path: Path to Transformer model directory (auto = automatic detection)
|
||
scheduler_path: Path to Scheduler config directory (auto = automatic detection)
|
||
|
||
Returns:
|
||
Processed video frames [B, H, W, C] in [0, 1]
|
||
"""
|
||
|
||
# Load models with auto-download support
|
||
self.load_models(vae_path, transformer_path, scheduler_path, auto_download)
|
||
|
||
print("🚀 Running BMO MiniMax-Remover")
|
||
print("=" * 50)
|
||
|
||
# Convert ComfyUI tensors to the format expected by MiniMax
|
||
batch_size, height, width, channels = images.shape
|
||
num_frames = batch_size
|
||
|
||
print(f"📐 Input: {images.shape} frames, {masks.shape} masks")
|
||
print(f"🎯 Parameters: steps={num_inference_steps}, iterations={iterations}, seed={seed}")
|
||
print(f"🤖 Auto-download: {'enabled' if auto_download else 'disabled'}")
|
||
|
||
# Show resolved paths
|
||
if vae_path == "auto" or transformer_path == "auto" or scheduler_path == "auto":
|
||
print(f"🗂️ Using auto-detected model paths")
|
||
else:
|
||
print(f"🗂️ Model paths: VAE={vae_path}, Transformer={transformer_path}, Scheduler={scheduler_path}")
|
||
|
||
# Prepare images: ComfyUI [B, H, W, C] -> MiniMax [F, H, W, C] -> [-1, 1]
|
||
images_np = images.detach().cpu().numpy() # [B, H, W, C] in [0, 1]
|
||
images_minimax = images_np * 2.0 - 1.0 # Convert to [-1, 1] for MiniMax
|
||
|
||
# Prepare masks: ComfyUI [B, H, W] -> MiniMax [F, H, W, C]
|
||
if len(masks.shape) == 3: # [B, H, W]
|
||
masks_np = masks.detach().cpu().numpy()
|
||
masks_minimax = np.expand_dims(masks_np, axis=-1) # [F, H, W, 1]
|
||
else: # [B, H, W, C]
|
||
masks_minimax = masks.detach().cpu().numpy()
|
||
|
||
print(f"🔄 Converted to MiniMax format:")
|
||
print(f" Images: {images_minimax.shape} [{images_minimax.min():.3f}, {images_minimax.max():.3f}]")
|
||
print(f" Masks: {masks_minimax.shape} [{masks_minimax.min():.3f}, {masks_minimax.max():.3f}]")
|
||
|
||
# Convert to tensors
|
||
images_tensor = torch.from_numpy(images_minimax).float()
|
||
masks_tensor = torch.from_numpy(masks_minimax).float()
|
||
|
||
# Set up generator for reproducible results
|
||
generator = torch.Generator(device=self.device).manual_seed(seed)
|
||
|
||
# Run the BMO pipeline
|
||
print(f"🔥 Processing with BMO pipeline...")
|
||
|
||
try:
|
||
with torch.no_grad():
|
||
result = self.pipe(
|
||
images=images_tensor,
|
||
masks=masks_tensor,
|
||
num_frames=num_frames,
|
||
height=height,
|
||
width=width,
|
||
num_inference_steps=num_inference_steps,
|
||
iterations=iterations,
|
||
generator=generator,
|
||
output_type="np" # Get numpy output
|
||
)
|
||
|
||
# Extract frames
|
||
output_frames = result.frames
|
||
print(f"✅ Pipeline completed!")
|
||
print(f"📊 Output: {output_frames.shape} [{output_frames.min():.3f}, {output_frames.max():.3f}]")
|
||
|
||
# Convert back to ComfyUI format [B, H, W, C]
|
||
if len(output_frames.shape) == 5: # [1, F, H, W, C]
|
||
output_frames = output_frames[0] # Remove batch dimension
|
||
|
||
# Ensure we have the right number of frames
|
||
if output_frames.shape[0] != num_frames:
|
||
print(f"🔧 Adjusting frame count: {output_frames.shape[0]} -> {num_frames}")
|
||
if output_frames.shape[0] < num_frames:
|
||
# Pad by repeating last frame
|
||
last_frame = output_frames[-1:].repeat(num_frames - output_frames.shape[0], axis=0)
|
||
output_frames = np.concatenate([output_frames, last_frame], axis=0)
|
||
else:
|
||
# Truncate to desired number of frames
|
||
output_frames = output_frames[:num_frames]
|
||
|
||
# Ensure correct shape and range
|
||
output_frames = np.clip(output_frames, 0.0, 1.0)
|
||
|
||
print(f"📤 Final output: {output_frames.shape} [{output_frames.min():.3f}, {output_frames.max():.3f}]")
|
||
|
||
# Convert back to tensor for ComfyUI
|
||
result_tensor = torch.from_numpy(output_frames).float()
|
||
|
||
return (result_tensor,)
|
||
|
||
except Exception as e:
|
||
print(f"❌ Processing failed: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
# Return original images as fallback
|
||
return (images,)
|
||
|
||
def diagnose_inputs(self, images, masks):
|
||
"""
|
||
Diagnostic function to check input compatibility
|
||
Helps users identify potential dimension issues before processing
|
||
"""
|
||
print("🔍 DIAGNOSTIC MODE: Analyzing inputs for compatibility")
|
||
print("=" * 60)
|
||
|
||
# Analyze images
|
||
if images is None:
|
||
print("❌ ERROR: Images tensor is None")
|
||
return False
|
||
|
||
print(f"📊 Images analysis:")
|
||
print(f" Shape: {images.shape}")
|
||
print(f" Type: {type(images)}")
|
||
print(f" Data type: {images.dtype if hasattr(images, 'dtype') else 'N/A'}")
|
||
print(f" Range: [{images.min():.3f}, {images.max():.3f}]" if hasattr(images, 'min') else "")
|
||
|
||
if len(images.shape) != 4:
|
||
print(f"❌ ERROR: Expected 4D tensor [B, H, W, C], got {len(images.shape)}D")
|
||
return False
|
||
|
||
batch_size, height, width, channels = images.shape
|
||
print(f" Frames: {batch_size}")
|
||
print(f" Resolution: {height}x{width}")
|
||
print(f" Channels: {channels}")
|
||
|
||
# Analyze masks
|
||
if masks is None:
|
||
print("❌ ERROR: Masks tensor is None")
|
||
return False
|
||
|
||
print(f"\n Masks analysis:")
|
||
print(f" Shape: {masks.shape}")
|
||
print(f" Type: {type(masks)}")
|
||
print(f" Data type: {masks.dtype if hasattr(masks, 'dtype') else 'N/A'}")
|
||
print(f" Range: [{masks.min():.3f}, {masks.max():.3f}]" if hasattr(masks, 'min') else "")
|
||
|
||
# Check dimension compatibility
|
||
print(f"\n🔧 Compatibility analysis:")
|
||
|
||
# VAE spatial compatibility
|
||
vae_h = ((height + 7) // 8) * 8
|
||
vae_w = ((width + 7) // 8) * 8
|
||
if height == vae_h and width == vae_w:
|
||
print(f" ✅ Resolution {height}x{width} is VAE-compatible")
|
||
else:
|
||
print(f" ⚠️ Resolution {height}x{width} will be adjusted to {vae_h}x{vae_w} for VAE compatibility")
|
||
|
||
# Temporal compatibility
|
||
vae_scale_factor_temporal = 4 # Default value
|
||
temporal_latent_frames = (batch_size - 1) // vae_scale_factor_temporal + 1
|
||
print(f" 📊 Temporal: {batch_size} frames → {temporal_latent_frames} latent frames")
|
||
|
||
# Check common issues
|
||
issues = []
|
||
|
||
if channels != 3:
|
||
issues.append(f"Expected 3 channels (RGB), got {channels}")
|
||
|
||
if batch_size < 1:
|
||
issues.append(f"Invalid frame count: {batch_size}")
|
||
|
||
if len(masks.shape) not in [3, 4]:
|
||
issues.append(f"Masks should be 3D [F,H,W] or 4D [F,H,W,1], got {len(masks.shape)}D")
|
||
|
||
if len(masks.shape) == 3 and masks.shape != (batch_size, height, width):
|
||
issues.append(f"Mask shape {masks.shape} doesn't match image frames")
|
||
|
||
if len(masks.shape) == 4 and masks.shape != (batch_size, height, width, 1):
|
||
issues.append(f"Mask shape {masks.shape} doesn't match expected [F,H,W,1]")
|
||
|
||
# Report results
|
||
if issues:
|
||
print(f"\n❌ ISSUES DETECTED:")
|
||
for i, issue in enumerate(issues, 1):
|
||
print(f" {i}. {issue}")
|
||
return False
|
||
else:
|
||
print(f"\n✅ ALL CHECKS PASSED - Inputs are compatible!")
|
||
return True
|
||
|
||
|
||
# ComfyUI Node Mappings
|
||
NODE_CLASS_MAPPINGS = {
|
||
"MinimaxRemoverBMO": MinimaxRemoverBMONode,
|
||
}
|
||
|
||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||
"MinimaxRemoverBMO": "MiniMax-Remover (BMO)",
|
||
}
|
||
|
||
# Export for ComfyUI
|
||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] |