diff --git a/py/promptsFromJanusPro.py b/py/promptsFromJanusPro.py index 692fe72..9aca8de 100644 --- a/py/promptsFromJanusPro.py +++ b/py/promptsFromJanusPro.py @@ -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)