import os
import torch
import gradio as gr

from depth_anything_3.model import DepthAnything3


# ---------------------------------------------------------
#  CPU-safe DepthAnything3 loader
# ---------------------------------------------------------

def load_model_cpu(model_dir):
    print("🔄 Loading DepthAnything3 model on CPU...")
    model = DepthAnything3.from_pretrained(model_dir, config_name="config.json")
    model.to("cpu")
    model.eval()
    print("✅ Model ready on CPU")
    return model


MODEL_DIR = os.environ.get("DA3_MODEL_DIR", "depth-anything/DA3NESTED-GIANT-LARGE")
model = load_model_cpu(MODEL_DIR)


# ---------------------------------------------------------
#  CPU-safe inference (single image only)
# ---------------------------------------------------------

def run_depth(img):
    """
    CPU version of depth inference.
    No batching, no multiview, no GS, no reconstruction.
    """
    if img is None:
        return None

    with torch.no_grad():
        depth = model.infer_image(img, device="cpu")

    return depth


# ---------------------------------------------------------
#  Minimal Gradio UI (fast startup)
# ---------------------------------------------------------

title = "Depth Anything 3 — CPU Mode (Safe HF Version)"
description = """
This Hugging Face Space runs **DepthAnything3** in CPU-only mode.
Only single-image depth inference is enabled.  
All heavy multiview / GS / reconstruction features were removed so the Space can boot on CPU.
"""

demo = gr.Interface(
    fn=run_depth,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=gr.Image(label="Predicted Depth"),
    title=title,
    description=description,
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)