Update Prompts From Janus-Pro
This commit is contained in:
samuelxuliu
+316
-52
@@ -11,7 +11,9 @@ import io
|
||||
import json
|
||||
import hashlib
|
||||
import folder_paths
|
||||
from typing import Optional
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import Optional, Dict, List
|
||||
from PIL import Image
|
||||
import torch
|
||||
import numpy as np
|
||||
@@ -96,40 +98,234 @@ class S4PromptsFromJanusPro:
|
||||
FUNCTION = "analyze_image"
|
||||
CATEGORY = "💀PromptsO"
|
||||
DESCRIPTION = "Generate text prompts from images using local Janus-Pro models with dynamic model path detection."
|
||||
|
||||
def get_model_path(self, model_variant: str) -> str:
|
||||
"""Get the correct model path based on ComfyUI's model directories"""
|
||||
|
||||
|
||||
def check_dependencies(self) -> Dict[str, any]:
|
||||
"""Check if required dependencies are installed with version info"""
|
||||
dependencies = {
|
||||
'janus': {'installed': False, 'version': None, 'error': None},
|
||||
'transformers': {'installed': False, 'version': None, 'error': None},
|
||||
'torch': {'installed': False, 'version': None, 'error': None},
|
||||
'PIL': {'installed': False, 'version': None, 'error': None}
|
||||
}
|
||||
|
||||
# Try to get the models path from ComfyUI's folder_paths
|
||||
# Check janus
|
||||
try:
|
||||
from janus.models import MultiModalityCausalLM, VLChatProcessor
|
||||
import janus
|
||||
dependencies['janus']['installed'] = True
|
||||
dependencies['janus']['version'] = getattr(janus, '__version__', 'unknown')
|
||||
except ImportError as e:
|
||||
dependencies['janus']['error'] = str(e)
|
||||
except Exception as e:
|
||||
dependencies['janus']['error'] = f"Import error: {str(e)}"
|
||||
|
||||
# Check transformers
|
||||
try:
|
||||
import transformers
|
||||
dependencies['transformers']['installed'] = True
|
||||
dependencies['transformers']['version'] = transformers.__version__
|
||||
except ImportError as e:
|
||||
dependencies['transformers']['error'] = str(e)
|
||||
|
||||
# Check torch
|
||||
try:
|
||||
import torch
|
||||
dependencies['torch']['installed'] = True
|
||||
dependencies['torch']['version'] = torch.__version__
|
||||
except ImportError as e:
|
||||
dependencies['torch']['error'] = str(e)
|
||||
|
||||
# Check PIL
|
||||
try:
|
||||
from PIL import Image
|
||||
import PIL
|
||||
dependencies['PIL']['installed'] = True
|
||||
dependencies['PIL']['version'] = PIL.__version__
|
||||
except ImportError as e:
|
||||
dependencies['PIL']['error'] = str(e)
|
||||
|
||||
return dependencies
|
||||
|
||||
def install_janus_dependency(self) -> bool:
|
||||
"""Attempt to install janus dependency"""
|
||||
try:
|
||||
S4APILogger.info("JanusDependency", "Installing janus library...")
|
||||
|
||||
# Use the same Python executable that's running ComfyUI
|
||||
result = subprocess.run([
|
||||
sys.executable, "-m", "pip", "install",
|
||||
"git+https://github.com/deepseek-ai/Janus.git"
|
||||
], capture_output=True, text=True, timeout=300)
|
||||
|
||||
if result.returncode == 0:
|
||||
S4APILogger.info("JanusDependency", "Janus library installed successfully")
|
||||
return True
|
||||
else:
|
||||
S4APILogger.error("JanusDependency", f"Installation failed: {result.stderr}")
|
||||
return False
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
S4APILogger.error("JanusDependency", "Installation timeout (5 minutes)")
|
||||
return False
|
||||
except Exception as e:
|
||||
S4APILogger.error("JanusDependency", f"Installation error: {str(e)}")
|
||||
return False
|
||||
|
||||
def find_all_model_paths(self) -> List[str]:
|
||||
"""Find all possible model directory paths automatically"""
|
||||
potential_paths = []
|
||||
|
||||
# 1. Try ComfyUI's folder_paths (highest priority)
|
||||
try:
|
||||
models_dir = folder_paths.models_dir
|
||||
S4APILogger.debug("JanusProModel", f"ComfyUI models directory: {models_dir}")
|
||||
if models_dir and os.path.exists(models_dir):
|
||||
potential_paths.append(models_dir)
|
||||
S4APILogger.debug("JanusModelPaths", f"ComfyUI models_dir: {models_dir}")
|
||||
except:
|
||||
# Fallback: try common ComfyUI model paths
|
||||
possible_paths = [
|
||||
os.path.join(os.path.dirname(folder_paths.__file__), "..", "models"),
|
||||
os.path.join(os.path.dirname(os.path.dirname(__file__)), "..", "..", "models"),
|
||||
"D:\\AI\\Models", # User's specific path
|
||||
"C:\\iCloud_Drive\\AI\\Models" # Alternative example path
|
||||
]
|
||||
pass
|
||||
|
||||
# 2. Relative to current ComfyUI installation (high priority)
|
||||
try:
|
||||
# From custom_nodes back to ComfyUI root
|
||||
comfyui_root = os.path.dirname(os.path.dirname(os.path.dirname(__file__)))
|
||||
models_path = os.path.join(comfyui_root, "models")
|
||||
if os.path.exists(models_path):
|
||||
potential_paths.append(models_path)
|
||||
S4APILogger.debug("JanusModelPaths", f"ComfyUI root models: {models_path}")
|
||||
except:
|
||||
pass
|
||||
|
||||
# 3. Search all drives for common ComfyUI installations
|
||||
import string
|
||||
available_drives = ['%s:' % d for d in string.ascii_uppercase if os.path.exists('%s:' % d)]
|
||||
|
||||
# Most common installation patterns
|
||||
common_patterns = [
|
||||
# Direct ComfyUI installations
|
||||
'ComfyUI\\models',
|
||||
'ComfyUI-Desktop\\models',
|
||||
'ComfyUI-portable\\ComfyUI\\models',
|
||||
'ComfyUI_windows_portable\\ComfyUI\\models',
|
||||
|
||||
models_dir = None
|
||||
for path in possible_paths:
|
||||
if os.path.exists(path):
|
||||
models_dir = os.path.abspath(path)
|
||||
print(f" • Found models directory: {models_dir}")
|
||||
break
|
||||
# AI tool collections
|
||||
'AI\\ComfyUI\\models',
|
||||
'AI\\Models',
|
||||
'AI\\models\\ComfyUI',
|
||||
|
||||
# Specific versions
|
||||
'ComfyUI-akr-v1.4\\ComfyUI-akr-v1.4\\models',
|
||||
'ComfyUI-akr-v1.5\\ComfyUI-akr-v1.5\\models',
|
||||
|
||||
# Alternative locations
|
||||
'Models\\ComfyUI',
|
||||
'machine-learning\\models',
|
||||
'deep-learning\\models',
|
||||
'stable-diffusion\\models',
|
||||
'ComfyUI\\custom_nodes\\models'
|
||||
]
|
||||
|
||||
if not models_dir:
|
||||
raise ValueError("Could not locate ComfyUI models directory")
|
||||
for drive in available_drives:
|
||||
for pattern in common_patterns:
|
||||
full_path = os.path.join(drive, os.sep, pattern)
|
||||
if os.path.exists(full_path):
|
||||
potential_paths.append(full_path)
|
||||
S4APILogger.debug("JanusModelPaths", f"Found drive path: {full_path}")
|
||||
|
||||
# Build path to Janus-Pro model
|
||||
model_path = os.path.join(models_dir, "Janus-Pro", model_variant)
|
||||
# 4. Environment variable (user override)
|
||||
env_path = os.environ.get('COMFYUI_MODELS_PATH')
|
||||
if env_path and os.path.exists(env_path):
|
||||
potential_paths.insert(0, env_path) # Highest priority
|
||||
S4APILogger.info("JanusModelPaths", f"Using env COMFYUI_MODELS_PATH: {env_path}")
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
raise ValueError(f"Janus-Pro model not found at: {model_path}")
|
||||
# Remove duplicates while preserving order (first occurrence has priority)
|
||||
seen = set()
|
||||
unique_paths = []
|
||||
for path in potential_paths:
|
||||
normalized_path = os.path.normpath(path.lower())
|
||||
if normalized_path not in seen:
|
||||
seen.add(normalized_path)
|
||||
unique_paths.append(path)
|
||||
|
||||
return model_path
|
||||
S4APILogger.info("JanusModelPaths", f"Scanning {len(unique_paths)} directories for Janus-Pro models")
|
||||
|
||||
return unique_paths
|
||||
|
||||
def get_model_path(self, model_variant: str) -> str:
|
||||
"""Get the correct model path with intelligent path discovery"""
|
||||
|
||||
# Get all potential model directories
|
||||
potential_model_dirs = self.find_all_model_paths()
|
||||
|
||||
found_models = []
|
||||
for models_dir in potential_model_dirs:
|
||||
# Check for Janus-Pro directory
|
||||
janus_dir = os.path.join(models_dir, "Janus-Pro")
|
||||
if os.path.exists(janus_dir):
|
||||
# Check for specific model variant
|
||||
model_path = os.path.join(janus_dir, model_variant)
|
||||
if os.path.exists(model_path):
|
||||
# Verify model files are complete - check for different possible file combinations
|
||||
required_files_combinations = [
|
||||
['config.json', 'pytorch_model.bin', 'tokenizer.json'],
|
||||
['config.json', 'model.safetensors', 'tokenizer.json'],
|
||||
['config.json', 'pytorch_model.bin', 'tokenizer_config.json'],
|
||||
['config.json', 'model.safetensors', 'tokenizer_config.json']
|
||||
]
|
||||
|
||||
model_complete = False
|
||||
for required_files in required_files_combinations:
|
||||
files_exist = all(os.path.exists(os.path.join(model_path, f)) for f in required_files)
|
||||
if files_exist:
|
||||
model_complete = True
|
||||
S4APILogger.info("JanusProModel", f"Found complete model at: {model_path}")
|
||||
S4APILogger.debug("JanusProModel", f"Model files: {required_files}")
|
||||
return model_path
|
||||
|
||||
if not model_complete:
|
||||
# List actual files found for debugging
|
||||
actual_files = [f for f in os.listdir(model_path) if os.path.isfile(os.path.join(model_path, f))]
|
||||
S4APILogger.warning("JanusProModel", f"Incomplete model found at: {model_path}")
|
||||
S4APILogger.debug("JanusProModel", f"Found files: {actual_files[:10]}...") # Limit output
|
||||
found_models.append((model_path, False))
|
||||
else:
|
||||
S4APILogger.debug("JanusProModel", f"Model variant {model_variant} not found in {janus_dir}")
|
||||
|
||||
# Provide simple and direct error message
|
||||
error_msg = self.generate_simple_model_error(model_variant, potential_model_dirs, found_models)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
def generate_simple_model_error(self, model_variant: str, searched_paths: List[str], found_models: List) -> str:
|
||||
"""Generate simple and actionable error message"""
|
||||
|
||||
error_msg = f"❌ {model_variant} model not found!\n\n"
|
||||
|
||||
error_msg += "🔧 Quick Solution:\n"
|
||||
error_msg += f"1. Download the model:\n"
|
||||
error_msg += f" git clone https://huggingface.co/deepseek-ai/{model_variant}\n\n"
|
||||
|
||||
if searched_paths and os.path.exists(searched_paths[0]):
|
||||
target_path = os.path.join(searched_paths[0], "Janus-Pro", model_variant)
|
||||
error_msg += f"2. Place it in: {target_path}\n\n"
|
||||
else:
|
||||
error_msg += "2. Place it in your ComfyUI/models/Janus-Pro/ folder\n\n"
|
||||
|
||||
error_msg += "📝 Model structure should be:\n"
|
||||
error_msg += " models/Janus-Pro/Janus-Pro-1B/\n"
|
||||
error_msg += " ├── config.json\n"
|
||||
error_msg += " ├── pytorch_model.bin\n"
|
||||
error_msg += " └── tokenizer files...\n\n"
|
||||
|
||||
if found_models:
|
||||
error_msg += f"⚠️ Found incomplete model at: {found_models[0][0]}\n"
|
||||
error_msg += " Please re-download or check for missing files.\n\n"
|
||||
|
||||
if len(searched_paths) > 0:
|
||||
error_msg += f"🔍 Searched in {len(searched_paths)} directories.\n"
|
||||
error_msg += f" Primary location: {searched_paths[0]}\n"
|
||||
|
||||
return error_msg
|
||||
|
||||
|
||||
def tensor_to_pil(self, image_tensor) -> Image.Image:
|
||||
"""Convert ComfyUI image tensor to PIL Image"""
|
||||
@@ -214,7 +410,58 @@ class S4PromptsFromJanusPro:
|
||||
S4APILogger.warning("JanusProCache", f"Failed to save cache: {e}")
|
||||
|
||||
def load_janus_model(self, model_path: str):
|
||||
"""Load Janus-Pro model and processor"""
|
||||
"""Load Janus-Pro model and processor with dependency checking"""
|
||||
|
||||
# Check dependencies first
|
||||
deps = self.check_dependencies()
|
||||
missing_deps = [dep for dep, info in deps.items() if not info['installed']]
|
||||
|
||||
if missing_deps:
|
||||
# Log detailed dependency info for debugging
|
||||
for dep, info in deps.items():
|
||||
if info['installed']:
|
||||
S4APILogger.info("JanusDependency", f"{dep}: ✅ v{info['version']}")
|
||||
else:
|
||||
S4APILogger.warning("JanusDependency", f"{dep}: ❌ {info['error'] or 'Not found'}")
|
||||
|
||||
if 'janus' in missing_deps:
|
||||
S4APILogger.info("JanusDependency", "Janus library not found, attempting auto-install...")
|
||||
if self.install_janus_dependency():
|
||||
# Retry dependency check after installation
|
||||
deps = self.check_dependencies()
|
||||
missing_deps = [dep for dep, info in deps.items() if not info['installed']]
|
||||
else:
|
||||
raise RuntimeError(
|
||||
"❌ Janus library installation failed!\n\n"
|
||||
"🔧 Manual installation required:\n"
|
||||
"pip install git+https://github.com/deepseek-ai/Janus.git"
|
||||
)
|
||||
|
||||
if missing_deps:
|
||||
error_msg = f"❌ Missing required dependencies: {', '.join(missing_deps)}\n\n"
|
||||
error_msg += "🔧 Installation commands:\n\n"
|
||||
|
||||
for dep in missing_deps:
|
||||
if dep == 'janus':
|
||||
error_msg += "# Install Janus-Pro support\n"
|
||||
error_msg += "pip install git+https://github.com/deepseek-ai/Janus.git\n\n"
|
||||
elif dep == 'transformers':
|
||||
error_msg += "# Install Hugging Face Transformers\n"
|
||||
error_msg += "pip install transformers\n\n"
|
||||
elif dep == 'torch':
|
||||
error_msg += "# Install PyTorch (choose appropriate version)\n"
|
||||
error_msg += "pip install torch\n\n"
|
||||
elif dep == 'PIL':
|
||||
error_msg += "# Install Pillow for image processing\n"
|
||||
error_msg += "pip install Pillow\n\n"
|
||||
|
||||
# Add dependency error details
|
||||
error_msg += "📋 Detailed error information:\n"
|
||||
for dep in missing_deps:
|
||||
if deps[dep]['error']:
|
||||
error_msg += f" • {dep}: {deps[dep]['error']}\n"
|
||||
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
cache_key = ("janus", model_path)
|
||||
|
||||
@@ -227,54 +474,71 @@ class S4PromptsFromJanusPro:
|
||||
|
||||
# Check if model path exists
|
||||
if not os.path.exists(model_path):
|
||||
error_msg = (
|
||||
f"❌ Model path does not exist: {model_path}\n\n"
|
||||
"Please download the Janus-Pro model:\n\n"
|
||||
"1. Using git clone:\n"
|
||||
" git clone https://huggingface.co/deepseek-ai/Janus-Pro-1B\n"
|
||||
f" and move to {model_path}\n\n"
|
||||
"2. Or download manually and place in correct directory.\n"
|
||||
)
|
||||
raise RuntimeError(error_msg)
|
||||
# This should be caught by get_model_path now, but just in case
|
||||
raise RuntimeError(f"Model path does not exist: {model_path}")
|
||||
|
||||
try:
|
||||
# Try native Janus loader first
|
||||
# Import after dependency check
|
||||
from janus.models import MultiModalityCausalLM, VLChatProcessor
|
||||
|
||||
S4APILogger.debug("JanusProModel", "Loading with native Janus processor...")
|
||||
S4APILogger.info("JanusProModel", f"Loading Janus-Pro model from: {model_path}")
|
||||
|
||||
if processor is None:
|
||||
processor = VLChatProcessor.from_pretrained(model_path, trust_remote_code=True)
|
||||
self._GLOBAL_PROCESSOR_CACHE[cache_key] = processor
|
||||
S4APILogger.debug("JanusProModel", "Processor loaded successfully")
|
||||
|
||||
if model is None:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
|
||||
|
||||
# Choose optimal dtype based on device and capability
|
||||
if torch.cuda.is_available():
|
||||
# Use bfloat16 if supported, otherwise float16
|
||||
try:
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
except:
|
||||
dtype = torch.float16
|
||||
else:
|
||||
dtype = torch.float32
|
||||
|
||||
S4APILogger.info("JanusProModel", f"Loading model with device={device}, dtype={dtype}")
|
||||
|
||||
# Load model without dtype parameter (not supported in some versions)
|
||||
model = MultiModalityCausalLM.from_pretrained(
|
||||
model_path,
|
||||
dtype=dtype,
|
||||
trust_remote_code=True
|
||||
)
|
||||
model = model.to(device).eval()
|
||||
|
||||
# Apply dtype and device after loading
|
||||
try:
|
||||
model = model.to(device=device, dtype=dtype).eval()
|
||||
except Exception as dtype_error:
|
||||
S4APILogger.warning("JanusProModel", f"Failed to set dtype {dtype}, using default: {dtype_error}")
|
||||
model = model.to(device=device).eval()
|
||||
self._GLOBAL_MODEL_CACHE[cache_key] = model
|
||||
S4APILogger.info("JanusProModel", "Model loaded successfully")
|
||||
|
||||
return processor, model
|
||||
|
||||
except ImportError as e:
|
||||
raise RuntimeError(
|
||||
f"\u274c Failed to import Janus models: {str(e)}\n\n"
|
||||
"Please ensure janus library is properly installed:\n"
|
||||
"pip install git+https://github.com/deepseek-ai/Janus.git"
|
||||
)
|
||||
except Exception as e:
|
||||
# If native Janus fails, provide clear instructions
|
||||
# Enhanced error message with more context
|
||||
error_msg = (
|
||||
f"❌ Janus-Pro model loading failed: {str(e)}\n\n"
|
||||
"Possible solutions:\n\n"
|
||||
"1. Confirm model path is correct:\n"
|
||||
f" {model_path}\n\n"
|
||||
"2. Confirm model files are complete:\n"
|
||||
" - config.json\n"
|
||||
" - pytorch_model.bin or model.safetensors\n"
|
||||
" - tokenizer files\n\n"
|
||||
"3. Try re-downloading the model:\n"
|
||||
f"\u274c Janus-Pro model loading failed: {str(e)}\n\n"
|
||||
"\ud83d\udd0d Troubleshooting steps:\n\n"
|
||||
"1. \ud83d\udcc1 Verify model files:\n"
|
||||
f" Path: {model_path}\n"
|
||||
" Required: config.json, pytorch_model.bin, tokenizer files\n\n"
|
||||
"2. \ud83d\udd04 Try re-downloading the model:\n"
|
||||
" git clone https://huggingface.co/deepseek-ai/Janus-Pro-1B\n\n"
|
||||
"4. Check dependency versions:\n"
|
||||
"3. \ud83d\udcbe Check available memory:\n"
|
||||
" Model requires significant GPU/RAM\n\n"
|
||||
"4. \ud83d\udc1b Check dependencies:\n"
|
||||
" pip list | grep -E 'transformers|torch|janus'\n"
|
||||
)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
Reference in New Issue
Block a user