hf_deploy/launcher.py

import os
import sys
import shutil
from pathlib import Path
from huggingface_hub import snapshot_download
import importlib.util

from dotenv import load_dotenv
load_dotenv()

def setup_cache_directory():
    """Create a hidden cache directory '.cache'.

    If an existing 'cache' directory exists, migrate its contents into '.cache'.
    Set restrictive permissions (owner rwx) and on macOS set the Finder hidden flag
    unless the environment variable `HIDE_CACHE` is explicitly set to '0'.
    """
    hidden_cache = Path(".cache")
    public_cache = Path("cache")

    # If the hidden cache doesn't exist but a public one does, move it.
    try:
        if not hidden_cache.exists():
            if public_cache.exists():
                # Prefer move to preserve contents and metadata
                public_cache.rename(hidden_cache)
            else:
                hidden_cache.mkdir(exist_ok=True)
        else:
            # ensure it exists
            hidden_cache.mkdir(exist_ok=True)

        # Restrict permissions to owner only (rwx------)
        try:
            hidden_cache.chmod(0o700)
        except Exception:
            # chmod may fail on some filesystems or platforms; ignore
            pass

        # On macOS, optionally set Finder hidden flag for extra concealment
        if sys.platform == "darwin":
            hide_flag = os.environ.get("HIDE_CACHE", "1")
            if hide_flag != "0":
                try:
                    # 'chflags hidden <path>' will make the folder hidden in Finder
                    os.system(f"/usr/bin/chflags hidden {hidden_cache}")
                except Exception:
                    pass

        return hidden_cache
    except Exception:
        # Fallback: try to create a simple cache folder named 'cache'
        public_cache.mkdir(exist_ok=True)
        return public_cache

def create_symbolic_links(cache_dir):
    """Create symbolic links for Generated, static, and view_session directories.
    
    The strategy is to make sure we have real directories in cache and symbolic links from root to cache.
    This ensures the BytePlus app can access these directories even when running from cache.
    """
    root_dir = Path.cwd()
    directories_to_link = ["Generated", "static", "view_session"]
    
    for dir_name in directories_to_link:
        # Root directory path (should be symlink)
        root_link = root_dir / dir_name
        # Cache directory path (should be real directory)
        cache_target = cache_dir / dir_name
        
        try:
            # Ensure we have a real directory in cache
            if cache_target.is_symlink():
                # Remove the symlink in cache if it exists
                cache_target.unlink()
                print(f"Removed cache symlink: {cache_target}")
            
            if not cache_target.exists():
                cache_target.mkdir(parents=True, exist_ok=True)
                print(f"Created cache directory: {cache_target}")
            
            # Set proper permissions on cache directory (755)
            cache_target.chmod(0o755)
            print(f"Set permissions 755 on: {cache_target}")
            
            # Ensure root has symlink to cache (if it doesn't already exist)
            if not root_link.is_symlink():
                if root_link.exists():
                    # If it's a real directory, we might need to move content or just replace
                    if root_link.is_dir():
                        # Try to move content to cache first
                        try:
                            for item in root_link.iterdir():
                                dest = cache_target / item.name
                                if not dest.exists():
                                    shutil.move(str(item), str(dest))
                            shutil.rmtree(root_link)
                            print(f"Moved content from {root_link} to {cache_target}")
                        except Exception as e:
                            print(f"Could not move content: {e}")
                            shutil.rmtree(root_link)
                    else:
                        root_link.unlink()
                
                # Create the symlink from root to cache
                root_link.symlink_to(cache_target)
                print(f"Created root symlink: {root_link} -> {cache_target}")
            else:
                print(f"Root symlink already exists: {root_link}")
            
        except Exception as e:
            print(f"Error setting up links for {dir_name}: {e}")
            # Fallback: ensure at least the cache directory exists
            try:
                if not cache_target.exists():
                    cache_target.mkdir(parents=True, exist_ok=True)
                    cache_target.chmod(0o755)
                    print(f"Created fallback cache directory: {cache_target}")
            except Exception as fe:
                print(f"Fallback failed for {dir_name}: {fe}")

def ensure_directory_permissions():
    """Ensure proper permissions on key directories in the current working directory."""
    directories_to_fix = ["Generated", "static", "view_session"]
    
    for dir_name in directories_to_fix:
        dir_path = Path(dir_name)
        try:
            if not dir_path.exists():
                # Create directory if it doesn't exist
                dir_path.mkdir(parents=True, exist_ok=True)
                print(f"Created directory: {dir_path}")
            
            # Set directory permissions to 755 (owner rwx, group rx, others rx)
            dir_path.chmod(0o755)
            print(f"Set permissions 755 on directory: {dir_path}")
            
            # Set permissions on all subdirectories and files
            for item in dir_path.rglob('*'):
                try:
                    if item.is_dir():
                        item.chmod(0o755)
                    else:
                        item.chmod(0o644)  # Files: owner rw, group r, others r
                except Exception as e:
                    print(f"Warning: Could not set permissions on {item}: {e}")
                    
        except Exception as e:
            print(f"Error ensuring permissions for {dir_name}: {e}")

def download_space(cache_dir):
    token = os.environ.get("HF_TOKEN")
    repo_id = os.environ.get("REPO_ID")
    
    if not token or not repo_id:
        return False
    
    try:
        snapshot_download(
            repo_id=repo_id,
            repo_type="space",
            local_dir=cache_dir,
            token=token
        )
        return True
    except:
        return False

def load_app(cache_dir):
    sys.path.insert(0, str(cache_dir))
    app_path = cache_dir / "app.py"
    spec = importlib.util.spec_from_file_location("app", app_path)
    app = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(app)

    # First try the common 'demo' attribute used by many Gradio Spaces
    if hasattr(app, "demo"):
        return app.demo

    # Otherwise, search for any attribute with a callable 'launch' method
    for name in dir(app):
        try:
            attr = getattr(app, name)
        except Exception:
            continue
        # Skip classes (types) since their 'launch' will be an unbound function
        if isinstance(attr, type):
            continue
        if hasattr(attr, "launch") and callable(getattr(attr, "launch")):
            return attr

    # Next, accept top-level callables that return an object with a 'launch' method.
    # This covers Spaces that expose a factory function instead of a 'demo' object.
    factory_names = [
        "create_secure_interface",
        "create_app",
        "create_demo",
        "main",
        "generator",
    ]

    for name in factory_names:
        if hasattr(app, name):
            try:
                candidate = getattr(app, name)
                if callable(candidate):
                    created = candidate()
                    if hasattr(created, "launch") and callable(getattr(created, "launch")):
                        return created
                    # If the factory itself is a Gradio Interface (callable with launch), return it
                    if hasattr(candidate, "launch") and callable(getattr(candidate, "launch")):
                        return candidate
            except Exception:
                # ignore failures from calling the factory and continue searching
                pass

    # Also accept any top-level callable that when called returns an object with 'launch'
    for name in dir(app):
        if name.startswith("__"):
            continue
        try:
            attr = getattr(app, name)
        except Exception:
            continue
        if callable(attr):
            try:
                created = attr()
                if hasattr(created, "launch") and callable(getattr(created, "launch")):
                    return created
            except Exception:
                # if calling fails, skip
                continue

    # If nothing found, raise a helpful error describing the problem
    available = [n for n in dir(app) if not n.startswith("__")]
    raise AttributeError(
        "Could not find a demo application in the downloaded space."
        f" Looked for attribute 'demo' and any attribute with a callable 'launch'."
        f" Available top-level names: {available}"
    )

if __name__ == "__main__":
    print("Setting up BytePlus Image Generation Studio...")
    
    # First ensure proper permissions on existing directories
    ensure_directory_permissions()
    
    # Setup cache directory
    cache_dir = setup_cache_directory()
    print(f"Cache directory: {cache_dir}")
    
    # Download the space
    if download_space(cache_dir):
        print("Successfully downloaded BytePlus space")
        
        # Create symbolic links for directory access
        create_symbolic_links(cache_dir)
        
        # Load and launch the app
        try:
            demo = load_app(cache_dir)
            print("Launching BytePlus Image Generation Studio...")
            demo.launch()
        except Exception as e:
            print(f"Error launching app: {e}")
            sys.exit(1)
    else:
        print("Failed to download space. Please check your HF_TOKEN and REPO_ID environment variables.")
        sys.exit(1)