Skip to content

Feature: Multi-environment compatibility (Jupyter, Colab, RunPod, Cloud) #13

Description

@slahiri

Feature Request: Multi-Environment Compatibility

Ensure the extension works seamlessly across different environments: local ComfyUI, Jupyter Notebooks, Google Colab, and ComfyUI Cloud platforms.


1. Environment Detection

Automatically detect the runtime environment and adjust behavior accordingly.

class EnvironmentDetector:
    @staticmethod
    def detect():
        """Detect current runtime environment"""
        env = {
            'type': 'local',
            'is_colab': False,
            'is_jupyter': False,
            'is_cloud': False,
            'has_gdrive': False,
            'base_path': None
        }

        # Check for Google Colab
        try:
            import google.colab
            env['is_colab'] = True
            env['type'] = 'colab'
            # Check if GDrive is mounted
            if os.path.exists('/content/drive/MyDrive'):
                env['has_gdrive'] = True
        except ImportError:
            pass

        # Check for Jupyter Notebook
        try:
            from IPython import get_ipython
            if get_ipython() is not None:
                env['is_jupyter'] = True
                if not env['is_colab']:
                    env['type'] = 'jupyter'
        except:
            pass

        # Check for cloud platforms
        cloud_indicators = [
            ('RUNPOD_POD_ID', 'runpod'),
            ('VAST_CONTAINERLABEL', 'vast'),
            ('PAPERSPACE_GRADIENT', 'paperspace'),
            ('LAMBDA_LABS', 'lambda'),
            ('THINKDIFFUSION', 'thinkdiffusion'),
        ]
        for env_var, platform in cloud_indicators:
            if os.environ.get(env_var):
                env['is_cloud'] = True
                env['type'] = platform
                break

        return env

2. Dynamic Path Resolution

Resolve model paths based on environment.

def get_models_base_path(env=None):
    """Get the appropriate models directory for current environment"""
    if env is None:
        env = EnvironmentDetector.detect()

    # Google Colab
    if env['is_colab']:
        save_to_gdrive = os.environ.get('SAVE_TO_GDRIVE', 'false').lower() == 'true'
        if save_to_gdrive and env['has_gdrive']:
            gdrive_base = os.environ.get('GDRIVE_COMFYUI_PATH', '/content/drive/MyDrive/ComfyUI')
            return os.path.join(gdrive_base, 'models')
        return '/content/ComfyUI/models'

    # RunPod
    if env['type'] == 'runpod':
        # RunPod typically uses /workspace
        workspace = os.environ.get('RUNPOD_WORKSPACE', '/workspace')
        return os.path.join(workspace, 'ComfyUI', 'models')

    # Vast.ai
    if env['type'] == 'vast':
        return '/workspace/ComfyUI/models'

    # ThinkDiffusion
    if env['type'] == 'thinkdiffusion':
        return '/home/user/ComfyUI/models'

    # Local / Default - use ComfyUI's folder_paths
    try:
        import folder_paths
        return folder_paths.models_dir
    except:
        return os.path.join(os.getcwd(), 'models')

3. Persistent Storage Detection

Detect if storage persists across sessions (important for cloud).

def get_persistent_storage_path(env=None):
    """Get path to persistent storage if available"""
    if env is None:
        env = EnvironmentDetector.detect()

    # Colab with GDrive
    if env['is_colab'] and env['has_gdrive']:
        return '/content/drive/MyDrive'

    # RunPod with network volume
    if env['type'] == 'runpod':
        network_vol = os.environ.get('RUNPOD_NETWORK_VOLUME')
        if network_vol and os.path.exists(network_vol):
            return network_vol

    # Vast.ai persistent storage
    if env['type'] == 'vast':
        if os.path.exists('/workspace/persistent'):
            return '/workspace/persistent'

    # Local - always persistent
    if env['type'] == 'local':
        return os.path.expanduser('~')

    return None

def warn_if_not_persistent():
    """Warn user if downloads won't persist"""
    persistent_path = get_persistent_storage_path()
    if persistent_path is None:
        return {
            'warning': True,
            'message': 'Downloads may not persist after session ends. Consider mounting persistent storage.'
        }
    return {'warning': False}

4. Cloud-Specific Optimizations

class CloudOptimizer:
    @staticmethod
    def get_download_settings(env=None):
        """Get optimized download settings for environment"""
        if env is None:
            env = EnvironmentDetector.detect()

        settings = {
            'chunk_size': 1024 * 1024,  # 1MB default
            'max_parallel': 3,
            'use_aria2': False,
            'show_progress': True
        }

        # Cloud environments often have faster networks
        if env['is_cloud']:
            settings['chunk_size'] = 8 * 1024 * 1024  # 8MB chunks
            settings['max_parallel'] = 5

        # Colab has good bandwidth
        if env['is_colab']:
            settings['chunk_size'] = 4 * 1024 * 1024  # 4MB chunks
            settings['max_parallel'] = 4

        # Jupyter notebook - show inline progress
        if env['is_jupyter']:
            settings['show_progress'] = 'notebook'  # Use tqdm.notebook

        return settings

5. Environment Info in UI

Show current environment info in the extension UI.

@routes.get("/workflow-models/environment")
async def get_environment_info(request):
    """Return current environment information"""
    env = EnvironmentDetector.detect()
    persistent = get_persistent_storage_path()

    return web.json_response({
        'environment': env['type'],
        'is_cloud': env['is_cloud'],
        'is_colab': env['is_colab'],
        'is_jupyter': env['is_jupyter'],
        'has_persistent_storage': persistent is not None,
        'persistent_path': persistent,
        'models_path': get_models_base_path(env),
        'warning': warn_if_not_persistent()
    })

Frontend display:

Environment: Google Colab
Storage: Google Drive (persistent)
Models Path: /content/drive/MyDrive/ComfyUI/models

6. Supported Platforms

Platform Detection Persistent Storage Notes
Local ComfyUI Default ✅ Always Standard installation
Google Colab google.colab ✅ GDrive Mount drive first
Jupyter Notebook IPython ✅ Local Use tqdm.notebook
RunPod RUNPOD_POD_ID ⚠️ Network Vol Mount network volume
Vast.ai VAST_* ⚠️ /workspace Check persistence
Paperspace PAPERSPACE_* ✅ /storage Gradient notebooks
ThinkDiffusion Path check ✅ Yes Managed platform
Lambda Labs LAMBDA_* ⚠️ Depends Check instance type

Implementation Priority

  1. High: Environment detection - Core functionality
  2. High: Dynamic path resolution - Required for cloud
  3. Medium: Persistent storage warning - User experience
  4. Medium: Cloud optimizations - Performance
  5. Low: UI environment display - Nice to have

Reference

Inspired by cloud compatibility needs from:

  • Google Colab users
  • RunPod/Vast.ai deployments
  • ThinkDiffusion platform

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions