app.py

import os
import sys
import subprocess

# Disable SSR mode - it causes issues with Zero GPU
os.environ["GRADIO_SSR_MODE"] = "false"

# Check if we're running on HF Spaces with Zero GPU
# HF Spaces sets SPACE_ID environment variable
IS_HF_SPACE = bool(os.environ.get("SPACE_ID"))
IS_ZERO_GPU = IS_HF_SPACE or os.environ.get("ZERO_GPU_MODE", "0") == "1" or os.environ.get("SPACES_ZERO_GPU", "0") == "1"

if IS_HF_SPACE:
    print(f"🤗 Running on Hugging Face Space: {os.environ.get('SPACE_ID')}")

# Try to import spaces for Zero GPU support
# Note: spaces is provided by HF Spaces environment, not from pip
try:
    import spaces
    HAS_SPACES = True
    print("✅ Spaces module available - Zero GPU support enabled")
except ImportError:
    HAS_SPACES = False
    spaces = None
    print("ℹ️ Spaces module not available - running without Zero GPU")

import gradio as gr


def _launch():
    # Add CropFormer and DeepLab to PYTHONPATH (needed when using prebuilt detectron2)
    # Prebuilt detectron2 doesn't include projects/, so we need to add them manually
    repo_root = os.path.dirname(os.path.abspath(__file__))
    projects_root = os.path.join(repo_root, "MaskClustering/third_party/detectron2/projects")

    # Add CropFormer (main project we use)
    cropformer_path = os.path.join(projects_root, "CropFormer")
    if cropformer_path not in sys.path:
        sys.path.insert(0, cropformer_path)

    # Add DeepLab (required by CropFormer for add_deeplab_config)
    deeplab_path = os.path.join(projects_root, "DeepLab")
    if deeplab_path not in sys.path:
        sys.path.insert(0, deeplab_path)

    # Set environment variables for Zero GPU optimization
    if IS_ZERO_GPU:
        print("🚀 Detected Zero GPU environment, optimizing settings...")
        os.environ["ZERO_GPU_MODE"] = "1"
        os.environ["MAX_GPU_MEMORY"] = "15GB"
        os.environ["BATCH_SIZE"] = "2"
        os.environ["MAX_IMAGES"] = "20"  # Limit max images for Zero GPU
        os.environ["MAX_RESOLUTION"] = "512"  # Lower resolution for memory optimization

    # Check if we need to install detectron2 and dependencies (first run only)
    # If using Dockerfile, DETECTRON2_INSTALLED is set and we skip compilation
    if not os.environ.get("DETECTRON2_INSTALLED"):
        print("🔧 Installing CropFormer dependencies...")

        # Check if pytorch3d is already installed (from Dockerfile prebuild)
        pytorch3d_installed = False
        try:
            import pytorch3d
            pytorch3d_installed = True
            print(f"✅ pytorch3d already installed (version: {getattr(pytorch3d, '__version__', 'unknown')})")
        except ImportError:
            pass

        try:
            # Install setuptools first (needed for mmcv build)
            subprocess.check_call([
                sys.executable, "-m", "pip", "install", "setuptools<81"
            ])
            # Install CropFormer/Mask2Former dependencies
            # Use --no-build-isolation to use system setuptools with pkg_resources
            subprocess.check_call([
                sys.executable, "-m", "pip", "install", "--no-build-isolation",
                "mmcv>=1.4.0,<2.0.0"
            ])

            # Only install pytorch3d if not already installed
            if not pytorch3d_installed:
                # Try prebuilt wheel first (much faster than compiling)
                # Using third-party wheels from https://github.com/MiroPsota/torch_packages_builder
                print("🔧 Installing pytorch3d from prebuilt wheel...")
                try:
                    subprocess.check_call([
                        sys.executable, "-m", "pip", "install",
                        "--extra-index-url", "https://miropsota.github.io/torch_packages_builder",
                        "pytorch3d"
                    ])
                    print("✅ pytorch3d installed from prebuilt wheel!")
                except subprocess.CalledProcessError:
                    # Fallback to compiling from source if prebuilt wheel not available
                    print("⚠️ Prebuilt wheel not found, compiling from source...")
                    cuda_env = {
                        **os.environ,
                        "FORCE_CUDA": "1",
                        "TORCH_CUDA_ARCH_LIST": "7.5;8.0;8.6;8.9;9.0",
                    }
                    subprocess.check_call([
                        sys.executable, "-m", "pip", "install", "--no-build-isolation",
                        "git+https://github.com/facebookresearch/pytorch3d.git"
                    ], env=cuda_env)
            print("✅ CropFormer dependencies installed!")

            # Check if detectron2 is already installed (from Dockerfile prebuild)
            detectron2_installed = False
            try:
                import detectron2
                detectron2_installed = True
                print(f"✅ detectron2 already installed (version: {getattr(detectron2, '__version__', 'unknown')})")
            except ImportError:
                pass

            if not detectron2_installed:
                # Use prebuilt wheels from https://miropsota.github.io/torch_packages_builder
                # This is MUCH faster than compiling from source (~5s vs ~60s)
                print("🔧 Installing detectron2 from prebuilt wheel...")
                try:
                    subprocess.check_call([
                        sys.executable, "-m", "pip", "install",
                        "--extra-index-url", "https://miropsota.github.io/torch_packages_builder",
                        "detectron2"
                    ])
                    print("✅ detectron2 installed from prebuilt wheel!")
                except subprocess.CalledProcessError:
                    # Fallback to compiling from vendored source
                    print("⚠️ Prebuilt wheel not found, compiling from vendored source...")
                    cuda_env = {
                        **os.environ,
                        "FORCE_CUDA": "1",
                        "TORCH_CUDA_ARCH_LIST": "7.5;8.0;8.6;8.9;9.0",
                    }
                    subprocess.check_call([
                        sys.executable, "-m", "pip", "install",
                        "--no-build-isolation", "-e", "MaskClustering/third_party/detectron2"
                    ], env=cuda_env)
                    print("✅ detectron2 compiled and installed!")

                # Install MSDeformAttn from prebuilt wheel
                print("🔧 Installing MultiScaleDeformableAttention from prebuilt wheel...")
                try:
                    subprocess.check_call([
                        sys.executable, "-m", "pip", "install",
                        "--extra-index-url", "https://miropsota.github.io/torch_packages_builder",
                        "MultiScaleDeformableAttention"
                    ])
                    print("✅ MultiScaleDeformableAttention installed from prebuilt wheel!")
                except subprocess.CalledProcessError:
                    # Fallback to compiling from source
                    print("⚠️ Prebuilt wheel not found, compiling MSDeformAttn from source...")
                    repo_root = os.path.dirname(os.path.abspath(__file__))
                    cropformer_root = os.path.join(repo_root, "MaskClustering/third_party/detectron2/projects/CropFormer")
                    ops_dir = os.path.join(cropformer_root, "mask2former/modeling/pixel_decoder/ops")
                    if 'cuda_env' not in locals():
                        cuda_env = {
                            **os.environ,
                            "FORCE_CUDA": "1",
                            "TORCH_CUDA_ARCH_LIST": "7.5;8.0;8.6;8.9;9.0",
                        }
                    subprocess.check_call(
                        [sys.executable, "-m", "pip", "install", "-e", ".", "--no-build-isolation"],
                        cwd=ops_dir,
                        env=cuda_env
                    )
                    print("✅ MSDeformAttn compiled and installed!")

                # Compile entity_api PythonAPI (still needed, but fast)
                repo_root = os.path.dirname(os.path.abspath(__file__))
                cropformer_root = os.path.join(repo_root, "MaskClustering/third_party/detectron2/projects/CropFormer")
                print("🔧 Compiling entity_api PythonAPI...")
                entity_api_dir = os.path.join(cropformer_root, "entity_api/PythonAPI")
                if os.path.exists(entity_api_dir):
                    try:
                        subprocess.check_call(["make"], cwd=entity_api_dir, shell=True)
                        print("✅ entity_api compiled successfully!")
                    except subprocess.CalledProcessError as e:
                        print(f"⚠️  entity_api compilation failed (non-critical): {e}")
                else:
                    print(f"⚠️  entity_api directory not found: {entity_api_dir}")
            
            print("🔄 Restarting to load detectron2...")
            
            # Set environment variable to avoid reinstalling on restart
            os.environ["DETECTRON2_INSTALLED"] = "1"
            
            # Restart the script
            os.execv(sys.executable, [sys.executable] + sys.argv)
        except subprocess.CalledProcessError as e:
            print(f"❌ Failed to install dependencies: {e}")
            raise
    else:
        # Second run: verify imports work
        print("🔍 Verifying detectron2 and mmcv imports...")
        try:
            import detectron2
            import mmcv
            print(f"✅ detectron2 available (version: {getattr(detectron2, '__version__', 'unknown')})")
            print(f"✅ mmcv available (version: {getattr(mmcv, '__version__', 'unknown')})")
        except ImportError as e:
            print(f"❌ Required module cannot be imported: {e}")
            print(f"   sys.path: {sys.path}")
            raise
    # HF Spaces/Gradio sometimes calls /api_info regardless of `show_api=False`.
    # Some gradio_client versions crash when JSON schema uses boolean `additionalProperties`.
    # Patch defensively to avoid bringing down the whole app.
    try:
        import gradio_client.utils as _gcu

        if hasattr(_gcu, "_json_schema_to_python_type"):
            _orig = _gcu._json_schema_to_python_type

            def _json_schema_to_python_type_patched(schema, defs=None):
                if isinstance(schema, bool):
                    return "Any"
                return _orig(schema, defs)

            _gcu._json_schema_to_python_type = _json_schema_to_python_type_patched
    except Exception:
        pass

    # HF Spaces expects the app to listen on 0.0.0.0:7860 (PORT may be provided).
    import mvp

    port = int(os.getenv("PORT", "7860"))
    # `mvp` defines `demo` (gr.Blocks). We launch it here instead of inside `mvp.py`.
    mvp.demo.queue(max_size=20).launch(
        server_name="0.0.0.0",
        server_port=port,
        show_error=True,
        share=False,
        show_api=False,
    )


if __name__ == "__main__":
    _launch()