manav0506's picture
Sync deps: Dockerfile system deps + ffmpeg, single pip flow
c179ef4
Raw
History Blame Contribute Delete
18.2 kB
"""
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()