koirpraw-ComfyUI-GCP-storage-decouple/folder_paths_patch.py

"""
ComfyUI GCP Storage Integration Patch for folder_paths.py

This module patches the core ComfyUI folder_paths functionality to work seamlessly
with GCP Cloud Storage while maintaining backward compatibility.
"""

import os
import sys
import importlib.util

# Add the current directory to Python path for imports
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

try:
    from gcp_storage_manager import storage_manager, gcp_open, gcp_copy, gcp_exists
    GCP_INTEGRATION_AVAILABLE = True
except ImportError:
    print("Warning: GCP Storage Manager not available")
    GCP_INTEGRATION_AVAILABLE = False

# Original folder_paths module
import folder_paths

# Store original functions before patching
if not hasattr(folder_paths, '_original_get_filename_list'):
    folder_paths._original_get_filename_list = folder_paths.get_filename_list

if not hasattr(folder_paths, '_original_add_model_folder_path'):
    folder_paths._original_add_model_folder_path = folder_paths.add_model_folder_path

if not hasattr(folder_paths, '_original_get_save_image_path'):
    folder_paths._original_get_save_image_path = folder_paths.get_save_image_path

def patched_get_filename_list(folder_name):
    """
    Enhanced get_filename_list that checks both local files and GCP Storage
    """
    if not GCP_INTEGRATION_AVAILABLE or not storage_manager.enabled:
        # Fall back to original behavior
        if hasattr(folder_paths, '_original_get_filename_list'):
            return folder_paths._original_get_filename_list(folder_name)
        else:
            # Implement basic functionality if original not available
            folder_paths_list, extensions = folder_paths.folder_names_and_paths.get(folder_name, ([], set()))
            files = []
            for folder_path in folder_paths_list:
                if os.path.exists(folder_path):
                    for file in os.listdir(folder_path):
                        if any(file.lower().endswith(ext) for ext in extensions):
                            files.append(file)
            return files
    
    # Get local files first
    local_files = set()
    folder_paths_list, extensions = folder_paths.folder_names_and_paths.get(folder_name, ([], set()))
    
    for folder_path in folder_paths_list:
        if os.path.exists(folder_path):
            for file in os.listdir(folder_path):
                if any(file.lower().endswith(ext) for ext in extensions):
                    local_files.add(file)
    
    # Get GCP files
    gcp_files = set()
    for folder_path in folder_paths_list:
        gcs_prefix = storage_manager.get_gcs_path(folder_path)
        if gcs_prefix:
            try:
                blobs = storage_manager.bucket.list_blobs(prefix=gcs_prefix)
                for blob in blobs:
                    if blob.name.endswith('/'):  # Skip directories
                        continue
                    
                    filename = os.path.basename(blob.name)
                    if any(filename.lower().endswith(ext) for ext in extensions):
                        gcp_files.add(filename)
            except Exception as e:
                print(f"Warning: Could not list GCP files for {folder_name}: {e}")
    
    # Combine and return unique files
    all_files = sorted(local_files.union(gcp_files))
    return all_files

def patched_get_full_path(folder_name, filename):
    """
    Enhanced get_full_path that downloads from GCP if needed
    """
    folder_paths_list, _ = folder_paths.folder_names_and_paths.get(folder_name, ([], set()))
    
    for folder_path in folder_paths_list:
        full_path = os.path.join(folder_path, filename)
        
        # Check if file exists locally
        if os.path.exists(full_path):
            return full_path
        
        # If GCP integration is available, try to download
        if GCP_INTEGRATION_AVAILABLE and storage_manager.enabled:
            downloaded_path = storage_manager.download_if_needed(full_path)
            if downloaded_path != full_path and os.path.exists(downloaded_path):
                return downloaded_path
    
    # If not found anywhere, return the first folder path + filename
    if folder_paths_list:
        return os.path.join(folder_paths_list[0], filename)
    
    return None

def patched_get_save_image_path(filename_prefix, output_dir, image_width=0, image_height=0):
    """
    Enhanced save image path that handles GCP upload after saving
    """
    # Use original function to get the path
    original_func = getattr(folder_paths, '_original_get_save_image_path', None)
    if original_func:
        full_filename, filename, counter, subfolder, filename_prefix = original_func(
            filename_prefix, output_dir, image_width, image_height
        )
    else:
        # Fallback implementation
        counter = 1
        while True:
            filename = f"{filename_prefix}_{counter:05d}.png"
            full_filename = os.path.join(output_dir, filename)
            if not os.path.exists(full_filename):
                break
            counter += 1
        subfolder = ""
    
    # If GCP integration is available, ensure the file will be uploaded after save
    if GCP_INTEGRATION_AVAILABLE and storage_manager.enabled:
        # The actual upload will happen when the file is written
        pass
    
    return full_filename, filename, counter, subfolder, filename_prefix

def patch_folder_paths():
    """Apply patches to folder_paths module"""
    if not GCP_INTEGRATION_AVAILABLE:
        print("GCP Storage integration not available, skipping folder_paths patches")
        return False
    
    try:
        # Patch functions
        if hasattr(folder_paths, 'get_filename_list'):
            folder_paths.get_filename_list = patched_get_filename_list
        
        if hasattr(folder_paths, 'get_full_path'):
            folder_paths.get_full_path = patched_get_full_path
        else:
            # Add the function if it doesn't exist
            setattr(folder_paths, 'get_full_path', patched_get_full_path)
        
        if hasattr(folder_paths, 'get_save_image_path'):
            folder_paths.get_save_image_path = patched_get_save_image_path
        
        print("✅ Successfully patched folder_paths for GCP Storage integration")
        return True
        
    except Exception as e:
        print(f"❌ Failed to patch folder_paths: {e}")
        return False

def unpatch_folder_paths():
    """Remove patches from folder_paths module"""
    try:
        if hasattr(folder_paths, '_original_get_filename_list'):
            folder_paths.get_filename_list = folder_paths._original_get_filename_list
        
        if hasattr(folder_paths, '_original_get_save_image_path'):
            folder_paths.get_save_image_path = folder_paths._original_get_save_image_path
        
        print("✅ Unpatched folder_paths")
        return True
        
    except Exception as e:
        print(f"❌ Failed to unpatch folder_paths: {e}")
        return False

# Auto-patch on import if enabled
if GCP_INTEGRATION_AVAILABLE and storage_manager.enabled:
    patch_folder_paths()