app.py

"""
Evoneural MVP - Local 3D Mesh + Skybox Generation
Run: streamlit run app.py
Open: http://localhost:8501
"""

import os
import sys
from pathlib import Path

# Ensure project root is on path
ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import streamlit as st

OUTPUTS = ROOT / "outputs"
OUTPUTS.mkdir(exist_ok=True)

# Hugging Face Space: HF sets SPACE_ID when running in a Space
IS_HF_SPACE = bool(os.environ.get("SPACE_ID") or os.environ.get("SPACE_REPO_ID"))


def main() -> None:
    st.set_page_config(
        page_title="Evoneural MVP - Mesh & Skybox",
        page_icon="🎮",
        layout="wide",
    )
    if IS_HF_SPACE:
        st.title("EvoneuralIn3D – Mesh & Skybox")
        st.caption("Text → 3D mesh (TripoSR) and Text → 360° skybox (Stable Diffusion). Running on Hugging Face Space.")
    else:
        st.title("Evoneural MVP – Local Mesh & Skybox")
        st.caption("Text → 3D mesh (TripoSR) and Text → 360° skybox (Stable Diffusion). Runs on localhost.")

    # Sidebar: model setup (token + download)
    with st.sidebar:
        st.subheader("Stable Diffusion model")
        hf_token_env = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        if IS_HF_SPACE:
            if hf_token_env:
                st.success("HF_TOKEN is set (from Space secrets)")
            else:
                st.error("HF_TOKEN not set")
                st.caption("Add it in this Space: **Settings** → **Variables and secrets** → New secret: `HF_TOKEN`. Then restart the Space.")
        from scripts.skybox_generator import _default_local_weights_dir
        local_model = _default_local_weights_dir()
        if local_model:
            st.success("Local model: found")
            st.caption(os.path.basename(local_model))
        elif not IS_HF_SPACE:
            st.warning("No local model. Download below or need internet on first generate.")
        if not IS_HF_SPACE:
            hf_token = st.text_input(
                "Hugging Face token (optional, if behind firewall)",
                type="password",
                key="hf_token",
                placeholder="hf_...",
                help="Get a token at huggingface.co/settings/tokens",
            )
            if hf_token:
                os.environ["HF_TOKEN"] = hf_token
        if not IS_HF_SPACE and st.button("Download model (~4GB to ./weights/sd-v1-5)", key="btn_download"):
            with st.spinner("Downloading model... (may take several minutes)"):
                try:
                    from scripts.download_sd_model import download_sd_model
                    path = download_sd_model(token=hf_token or os.environ.get("HF_TOKEN"))
                    st.success(f"Model saved. Try generating a skybox.")
                    st.rerun()
                except Exception as e:
                    st.error(str(e))
                    st.caption("Set a Hugging Face token above if your network blocks Hugging Face.")
        # Environment check: TripoSR (useful in Space)
        from scripts.mesh_generator import find_triposr_root as _find_triposr
        triposr_ok = _find_triposr(str(ROOT)) is not None
        if triposr_ok:
            st.caption("TripoSR: ready")
        else:
            st.caption("TripoSR: not found (mesh tab will show instructions)")

    # In Space, Skybox and mesh (text→mesh and image→mesh) need Hub access: SD and TripoSR download models. Disable if no token to avoid 403.
    can_use_hub = bool(hf_token_env) or not IS_HF_SPACE
    if IS_HF_SPACE and not hf_token_env:
        st.warning("Set **HF_TOKEN** in Settings → Variables and secrets to enable Skybox and mesh generation (TripoSR also downloads its model from the Hub).")

    tab_mesh, tab_skybox = st.tabs(["🟦 Text → 3D Mesh", "🌅 Text → Skybox"])

    with tab_mesh:
        st.subheader("Generate 3D mesh from text")
        st.markdown(
            "Uses **Stable Diffusion** for text→image, then **TripoSR** for image→mesh. "
            "TripoSR repo must be cloned into `./TripoSR` (see README)."
        )
        prompt_mesh = st.text_input(
            "Prompt (e.g. for mesh)",
            value="A highly detailed, sci-fi mechanical drone with glowing blue accents.",
            key="mesh_prompt",
        )
        col1, col2 = st.columns(2)
        with col1:
            mesh_format = st.selectbox("Mesh format", ["glb", "obj"], key="mesh_fmt")
            seed_mesh = st.number_input("Seed (optional)", value=42, min_value=0, key="mesh_seed")
        with col2:
            use_image = st.checkbox("Use image (from outputs or upload)", value=False, key="use_img")

        with st.expander("Quality options (TripoSR)", expanded=True):
            mc_resolution = st.selectbox(
                "Mesh resolution",
                options=[256, 512],
                index=1,
                format_func=lambda x: f"{x} (faster)" if x == 256 else f"{x} (higher quality)",
                key="mesh_mc_res",
                help="Marching cubes grid. 512 gives finer, less blocky meshes.",
            )
            bake_texture = st.checkbox(
                "Bake texture atlas",
                value=True,
                key="mesh_bake_tex",
                help="Produces a texture map instead of vertex colors; usually looks cleaner.",
            )
            smooth_mesh = st.checkbox(
                "Smooth mesh",
                value=True,
                key="mesh_smooth",
                help="Light Laplacian smoothing to reduce blockiness.",
            )
            output_images = sorted(Path(OUTPUTS).glob("*.png"), key=lambda p: p.stat().st_mtime, reverse=True)
            output_images += sorted(Path(OUTPUTS).glob("*.jpg"), key=lambda p: p.stat().st_mtime, reverse=True)
            selected_from_outputs = None
            if use_image and output_images:
                opt_names = [f.name for f in output_images]
                k = "mesh_pick_output_img"
                if k in st.session_state and st.session_state[k] not in opt_names:
                    del st.session_state[k]
                pick = st.selectbox("Pick from outputs (e.g. previous mesh input)", ["(upload below)"] + opt_names, key=k)
                if pick and pick != "(upload below)":
                    selected_from_outputs = Path(OUTPUTS) / pick
            uploaded = st.file_uploader("Or upload image for mesh", type=["png", "jpg"], key="mesh_upload") if use_image else None

        image_path_to_use = None
        if use_image and (selected_from_outputs and selected_from_outputs.exists() or uploaded):
            image_path_to_use = str(selected_from_outputs) if (selected_from_outputs and selected_from_outputs.exists()) else "upload"

        if st.button("Generate mesh", key="btn_mesh", disabled=not can_use_hub):
            if not image_path_to_use and not prompt_mesh.strip():
                st.warning("Enter a prompt or choose/upload an image.")
            else:
                with st.spinner("Running pipeline..."):
                    try:
                        from scripts.mesh_generator import (
                            generate_mesh_from_image,
                            generate_mesh_from_text,
                            find_triposr_root,
                        )
                        triposr_root = find_triposr_root(str(ROOT))
                        if not triposr_root:
                            st.error(
                                "TripoSR not found. In this Space the Docker image should include it. "
                                "If you see this, rebuild the Space or check the Dockerfile."
                            )
                        elif image_path_to_use:
                            path = image_path_to_use
                            if path == "upload" and uploaded:
                                path = os.path.join(OUTPUTS, "uploaded_mesh_input.png")
                                with open(path, "wb") as f:
                                    f.write(uploaded.getvalue())
                            if path != "upload" and os.path.isfile(path):
                                import torch as _torch
                                _dev = "cuda:0" if _torch.cuda.is_available() else "cpu"
                                mesh_path, elapsed, msg = generate_mesh_from_image(
                                    path,
                                    output_dir=str(OUTPUTS / "mesh_run"),
                                    mesh_format=mesh_format,
                                    mc_resolution=mc_resolution,
                                    bake_texture=bake_texture,
                                    smooth_mesh=smooth_mesh,
                                    device=_dev,
                                )
                                if mesh_path:
                                    st.success(f"Done in {elapsed:.1f}s. {msg}")
                                    with open(mesh_path, "rb") as f:
                                        mesh_data = f.read()
                                    st.download_button("Download mesh", data=mesh_data, file_name=os.path.basename(mesh_path), key="dl_mesh_upload")
                                else:
                                    st.error(msg)
                            elif path == "upload":
                                st.warning("Upload an image first.")
                        else:
                            import torch as _torch
                            _dev = "cuda:0" if _torch.cuda.is_available() else "cpu"
                            mesh_path, elapsed, msg = generate_mesh_from_text(
                                prompt_mesh,
                                output_dir=str(OUTPUTS),
                                mesh_format=mesh_format,
                                seed=seed_mesh,
                                mc_resolution=mc_resolution,
                                bake_texture=bake_texture,
                                smooth_mesh=smooth_mesh,
                                device=_dev,
                            )
                            if mesh_path:
                                st.success(f"Done in {elapsed:.1f}s. {msg}")
                                with open(mesh_path, "rb") as f:
                                    mesh_data = f.read()
                                st.download_button("Download mesh", data=mesh_data, file_name=os.path.basename(mesh_path), key="dl_mesh")
                            else:
                                st.error(msg)
                    except Exception as e:
                        st.exception(e)

        # View 3D mesh (GLB): path, upload, or pick from outputs
        with st.expander("View 3D mesh", expanded=False):
            st.caption("Open a .glb file by path, upload, or pick from outputs. Drag to rotate, scroll to zoom.")
            from scripts.mesh_viewer import mesh_viewer_html
            import streamlit.components.v1 as components

            viewer_glb_path: str | None = None
            viewer_glb_bytes: bytes | None = None

            path_input = st.text_input(
                "Path to .glb file",
                value="",
                key="mesh_viewer_path",
                placeholder=r"e.g. C:\Users\...\Downloads\mesh (1).glb",
            )
            if path_input and Path(path_input.strip()).is_file():
                viewer_glb_path = path_input.strip()

            uploaded_glb = st.file_uploader("Or upload a .glb file", type=["glb"], key="mesh_viewer_upload")
            if uploaded_glb is not None:
                viewer_glb_bytes = uploaded_glb.getvalue()

            output_glbs = sorted(Path(OUTPUTS).rglob("*.glb"), key=lambda p: p.stat().st_mtime, reverse=True)
            if not viewer_glb_path and not viewer_glb_bytes and output_glbs:
                opt_names = [str(p.relative_to(OUTPUTS)) for p in output_glbs]
                k = "mesh_viewer_pick"
                if k in st.session_state and st.session_state[k] not in opt_names:
                    del st.session_state[k]
                picked = st.selectbox("Or pick from outputs", ["(none)"] + opt_names, key=k)
                if picked and picked != "(none)":
                    viewer_glb_path = str(OUTPUTS / picked)

            if viewer_glb_path or viewer_glb_bytes:
                html = mesh_viewer_html(glb_path=viewer_glb_path, glb_bytes=viewer_glb_bytes, height_px=480)
                components.html(html, height=500, scrolling=False)
            else:
                st.info("Enter a path to a .glb file, upload one, or generate a mesh above and pick it from outputs.")

    with tab_skybox:
        st.subheader("Generate 2:1 equirectangular skybox")
        st.markdown(
            "Uses **Stable Diffusion 2.1** at 2:1 aspect (e.g. 1024×512). "
            "Optional seamless check compares left/right edges."
        )
        prompt_sky = st.text_input(
            "Prompt (e.g. for skybox)",
            value="Cyberpunk city skyline at dusk, neon reflections, cinematic lighting.",
            key="sky_prompt",
        )
        col1, col2 = st.columns(2)
        with col1:
            width = st.selectbox("Width", [1024, 2048], key="sky_w")
            height = width // 2
            seed_sky = st.number_input("Seed (optional)", value=42, min_value=0, key="sky_seed")
        with col2:
            check_seamless = st.checkbox("Run seamless edge check", value=True, key="seamless")

        if st.button("Generate skybox", key="btn_sky", disabled=not can_use_hub):
            if not prompt_sky.strip():
                st.warning("Enter a prompt.")
            else:
                try:
                    from scripts.skybox_generator import generate_skybox
                    from scripts.check_seamless import check_seamless as run_seamless

                    progress_placeholder = st.empty()
                    status_placeholder = st.empty()
                    progress_placeholder.progress(0)
                    status_placeholder.caption("Loading model and starting generation…")

                    def on_step(step: int, total: int) -> None:
                        progress = min(step / total, 1.0)
                        progress_placeholder.progress(progress)
                        status_placeholder.caption(f"Step {min(step, total)} / {total}")

                    out_path, elapsed, vram_mb = generate_skybox(
                        prompt_sky,
                        output_dir=str(OUTPUTS),
                        width=width,
                        height=height,
                        seed=seed_sky,
                        progress_callback=on_step,
                    )
                    progress_placeholder.progress(1.0)
                    status_placeholder.caption("Done.")

                    st.success(f"Done in {elapsed:.1f}s. Peak VRAM: {vram_mb:.0f} MB")
                    st.image(out_path, use_container_width=True)
                    with open(out_path, "rb") as f:
                        skybox_data = f.read()
                    st.download_button("Download skybox", data=skybox_data, file_name=os.path.basename(out_path), key="dl_sky")

                    if check_seamless:
                        result = run_seamless(out_path)
                        st.info(result["message"])

                    st.session_state["last_skybox_path"] = str(Path(out_path).resolve())
                except Exception as e:
                    st.exception(e)

        # Show 360° viewer for last generated skybox (same session)
        if "last_skybox_path" in st.session_state:
            last_path = Path(st.session_state["last_skybox_path"]).resolve()
            if last_path.exists():
                with st.expander("View in 360°", expanded=False):
                    st.caption("Drag to look around, scroll to zoom. Fullscreen available in the viewer.")
                    from scripts.panorama_viewer import panorama_html
                    import streamlit.components.v1 as components
                    components.html(panorama_html(last_path, height_px=480), height=500, scrolling=False)

        # View existing image from outputs (or upload) in 360° – test without regenerating
        with st.expander("View existing image in 360°", expanded=False):
            st.caption("Pick an image from outputs or upload a 2:1 equirectangular image to test the viewer.")
            from scripts.panorama_viewer import panorama_html
            import streamlit.components.v1 as components

            output_files = sorted(Path(OUTPUTS).glob("*.png"), key=lambda p: p.stat().st_mtime, reverse=True)
            viewer_path = None
            option_names = [f.name for f in output_files]

            if output_files:
                key = "skybox_select_existing"
                if key in st.session_state and st.session_state[key] not in option_names:
                    del st.session_state[key]
                selected_name = st.selectbox(
                    "Choose image from outputs",
                    options=option_names,
                    key=key,
                )
                if selected_name:
                    viewer_path = Path(OUTPUTS) / selected_name

            uploaded = st.file_uploader("Or upload a 2:1 equirectangular image", type=["png", "jpg", "jpeg"], key="skybox_upload_360")
            if uploaded is not None:
                upload_path = OUTPUTS / "uploaded_360_view.png"
                upload_path.write_bytes(uploaded.getvalue())
                viewer_path = upload_path

            if viewer_path is not None and viewer_path.exists():
                components.html(panorama_html(Path(viewer_path).resolve(), height_px=480), height=500, scrolling=False)
            elif not output_files and uploaded is None:
                st.info("No skybox images in outputs yet. Generate one above or upload an image.")

    st.divider()
    st.caption("Evoneural AI – Local ML Deployment MVP. Models run locally (no API).")


if __name__ == "__main__":
    main()