hf_app.py

#!/usr/bin/env python3
"""
HuggingFace Spaces version of BytePlus Image Generation Studio
This version is optimized for deployment on Hugging Face Spaces
"""

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

from dotenv import load_dotenv
load_dotenv()

def setup_huggingface_environment():
    """Setup environment for HuggingFace Spaces deployment"""
    
    # Detect if we're running on HuggingFace Spaces
    is_hf_space = bool(os.environ.get("SPACE_ID"))
    
    if is_hf_space:
        print("🌐 Detected HuggingFace Space environment")
        
        # Create essential directories that must exist for the app to work
        essential_dirs = ["Generated", "static", "view_session"]
        
        for dir_name in essential_dirs:
            dir_path = Path(dir_name)
            if not dir_path.exists():
                dir_path.mkdir(exist_ok=True, mode=0o755)
                print(f"✅ Created directory: {dir_name}")
            else:
                # Ensure correct permissions
                try:
                    dir_path.chmod(0o755)
                except Exception:
                    pass  # Ignore permission errors on some filesystems
        
        # For HuggingFace Spaces, we need to work with the current directory
        # since we can't download to a hidden cache directory
        return Path(".")
    else:
        # Local development - use the cache system
        return setup_cache_directory()

def setup_cache_directory():
    """Create a hidden cache directory '.cache' for local development"""
    hidden_cache = Path(".cache")
    public_cache = Path("cache")

    try:
        if not hidden_cache.exists():
            if public_cache.exists():
                public_cache.rename(hidden_cache)
            else:
                hidden_cache.mkdir(exist_ok=True)
        else:
            hidden_cache.mkdir(exist_ok=True)

        try:
            hidden_cache.chmod(0o700)
        except Exception:
            pass

        if sys.platform == "darwin":
            hide_flag = os.environ.get("HIDE_CACHE", "1")
            if hide_flag != "0":
                try:
                    os.system(f"/usr/bin/chflags hidden {hidden_cache}")
                except Exception:
                    pass

        return hidden_cache
    except Exception:
        public_cache.mkdir(exist_ok=True)
        return public_cache

def download_space(cache_dir):
    """Download the HuggingFace Space if needed"""
    token = os.environ.get("HF_TOKEN")
    repo_id = os.environ.get("REPO_ID")
    
    is_hf_space = bool(os.environ.get("SPACE_ID"))
    
    if is_hf_space:
        # On HuggingFace Spaces, the code is already available
        # We just need to ensure the main app file exists
        app_path = cache_dir / "app.py"
        if app_path.exists():
            print("✅ App code is already available in HuggingFace Space")
            return True
        else:
            print("❌ App code not found in HuggingFace Space")
            return False
    
    # Local development - download from HuggingFace
    if not token or not repo_id:
        print("⚠️ HF_TOKEN or REPO_ID not set")
        return False
    
    try:
        snapshot_download(
            repo_id=repo_id,
            repo_type="space",
            local_dir=cache_dir,
            token=token
        )
        return True
    except Exception as e:
        print(f"❌ Failed to download space: {e}")
        return False

def load_app(cache_dir):
    """Load the BytePlus app"""
    sys.path.insert(0, str(cache_dir))
    
    # Look for the main app file
    app_files = ["app.py", "main.py", "byteplus_app.py"]
    
    for app_file in app_files:
        app_path = cache_dir / app_file
        if app_path.exists():
            try:
                spec = importlib.util.spec_from_file_location("app", app_path)
                app = importlib.util.module_from_spec(spec)
                spec.loader.exec_module(app)
                
                # Try different ways to get the demo
                if hasattr(app, "demo"):
                    return app.demo
                elif hasattr(app, "create_secure_interface"):
                    return app.create_secure_interface()
                elif hasattr(app, "main"):
                    return app.main()
                    
            except Exception as e:
                print(f"⚠️ Failed to load {app_file}: {e}")
                continue
    
    raise AttributeError("Could not find or load the main app")

if __name__ == "__main__":
    print("🚀 Starting BytePlus Image Generation Studio")
    print("🌐 HuggingFace Spaces optimized version")
    
    # Setup environment
    cache_dir = setup_huggingface_environment()
    
    # Download space if needed (for local development)
    if download_space(cache_dir):
        print("✅ App loaded successfully")
        
        # For HuggingFace Spaces, run the app directly
        is_hf_space = bool(os.environ.get("SPACE_ID"))
        
        if is_hf_space:
            print("🌐 Running in HuggingFace Space mode")
            # Import and run the app directly
            try:
                sys.path.insert(0, str(cache_dir))
                
                # Since we're in HF Spaces, the app code should be here
                # Check if there's a main app file in the current directory
                if Path("app.py").exists():
                    import app as main_app
                    if hasattr(main_app, "demo"):
                        main_app.demo.launch()
                    elif hasattr(main_app, "launch"):
                        main_app.launch()
                    else:
                        print("❌ Could not find launch method in app")
                else:
                    print("❌ app.py not found in HuggingFace Space")
                    
            except Exception as e:
                print(f"❌ Error running in HuggingFace Space: {e}")
        else:
            # Local development
            demo = load_app(cache_dir)
            demo.launch()
    else:
        print("❌ Failed to setup app environment")