evoneural/evoneuralIn3D-app
0
1"""2Evoneural MVP - Local 3D Mesh + Skybox Generation3Run: streamlit run app.py4Open: http://localhost:85015"""6 7import os8import sys9from pathlib import Path10 11# Ensure project root is on path12ROOT = Path(__file__).resolve().parent13if str(ROOT) not in sys.path:14 sys.path.insert(0, str(ROOT))15 16import streamlit as st17 18OUTPUTS = ROOT / "outputs"19OUTPUTS.mkdir(exist_ok=True)20 21# Hugging Face Space: HF sets SPACE_ID when running in a Space22IS_HF_SPACE = bool(os.environ.get("SPACE_ID") or os.environ.get("SPACE_REPO_ID"))23 24 25def main() -> None:26 st.set_page_config(27 page_title="Evoneural MVP - Mesh & Skybox",28 page_icon="🎮",29 layout="wide",30 )31 if IS_HF_SPACE:32 st.title("EvoneuralIn3D – Mesh & Skybox")33 st.caption("Text → 3D mesh (TripoSR) and Text → 360° skybox (Stable Diffusion). Running on Hugging Face Space.")34 else:35 st.title("Evoneural MVP – Local Mesh & Skybox")36 st.caption("Text → 3D mesh (TripoSR) and Text → 360° skybox (Stable Diffusion). Runs on localhost.")37 38 # Sidebar: model setup (token + download)39 with st.sidebar:40 st.subheader("Stable Diffusion model")41 hf_token_env = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")42 if IS_HF_SPACE:43 if hf_token_env:44 st.success("HF_TOKEN is set (from Space secrets)")45 else:46 st.error("HF_TOKEN not set")47 st.caption("Add it in this Space: **Settings** → **Variables and secrets** → New secret: `HF_TOKEN`. Then restart the Space.")48 from scripts.skybox_generator import _default_local_weights_dir49 local_model = _default_local_weights_dir()50 if local_model:51 st.success("Local model: found")52 st.caption(os.path.basename(local_model))53 elif not IS_HF_SPACE:54 st.warning("No local model. Download below or need internet on first generate.")55 if not IS_HF_SPACE:56 hf_token = st.text_input(57 "Hugging Face token (optional, if behind firewall)",58 type="password",59 key="hf_token",60 placeholder="hf_...",61 help="Get a token at huggingface.co/settings/tokens",62 )63 if hf_token:64 os.environ["HF_TOKEN"] = hf_token65 if not IS_HF_SPACE and st.button("Download model (~4GB to ./weights/sd-v1-5)", key="btn_download"):66 with st.spinner("Downloading model... (may take several minutes)"):67 try:68 from scripts.download_sd_model import download_sd_model69 path = download_sd_model(token=hf_token or os.environ.get("HF_TOKEN"))70 st.success(f"Model saved. Try generating a skybox.")71 st.rerun()72 except Exception as e:73 st.error(str(e))74 st.caption("Set a Hugging Face token above if your network blocks Hugging Face.")75 # Environment check: TripoSR (useful in Space)76 from scripts.mesh_generator import find_triposr_root as _find_triposr77 triposr_ok = _find_triposr(str(ROOT)) is not None78 if triposr_ok:79 st.caption("TripoSR: ready")80 else:81 st.caption("TripoSR: not found (mesh tab will show instructions)")82 83 # 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.84 can_use_hub = bool(hf_token_env) or not IS_HF_SPACE85 if IS_HF_SPACE and not hf_token_env:86 st.warning("Set **HF_TOKEN** in Settings → Variables and secrets to enable Skybox and mesh generation (TripoSR also downloads its model from the Hub).")87 88 tab_mesh, tab_skybox = st.tabs(["🟦 Text → 3D Mesh", "🌅 Text → Skybox"])89 90 with tab_mesh:91 st.subheader("Generate 3D mesh from text")92 st.markdown(93 "Uses **Stable Diffusion** for text→image, then **TripoSR** for image→mesh. "94 "TripoSR repo must be cloned into `./TripoSR` (see README)."95 )96 prompt_mesh = st.text_input(97 "Prompt (e.g. for mesh)",98 value="A highly detailed, sci-fi mechanical drone with glowing blue accents.",99 key="mesh_prompt",100 )101 col1, col2 = st.columns(2)102 with col1:103 mesh_format = st.selectbox("Mesh format", ["glb", "obj"], key="mesh_fmt")104 seed_mesh = st.number_input("Seed (optional)", value=42, min_value=0, key="mesh_seed")105 with col2:106 use_image = st.checkbox("Use image (from outputs or upload)", value=False, key="use_img")107 108 with st.expander("Quality options (TripoSR)", expanded=True):109 mc_resolution = st.selectbox(110 "Mesh resolution",111 options=[256, 512],112 index=1,113 format_func=lambda x: f"{x} (faster)" if x == 256 else f"{x} (higher quality)",114 key="mesh_mc_res",115 help="Marching cubes grid. 512 gives finer, less blocky meshes.",116 )117 bake_texture = st.checkbox(118 "Bake texture atlas",119 value=True,120 key="mesh_bake_tex",121 help="Produces a texture map instead of vertex colors; usually looks cleaner.",122 )123 smooth_mesh = st.checkbox(124 "Smooth mesh",125 value=True,126 key="mesh_smooth",127 help="Light Laplacian smoothing to reduce blockiness.",128 )129 output_images = sorted(Path(OUTPUTS).glob("*.png"), key=lambda p: p.stat().st_mtime, reverse=True)130 output_images += sorted(Path(OUTPUTS).glob("*.jpg"), key=lambda p: p.stat().st_mtime, reverse=True)131 selected_from_outputs = None132 if use_image and output_images:133 opt_names = [f.name for f in output_images]134 k = "mesh_pick_output_img"135 if k in st.session_state and st.session_state[k] not in opt_names:136 del st.session_state[k]137 pick = st.selectbox("Pick from outputs (e.g. previous mesh input)", ["(upload below)"] + opt_names, key=k)138 if pick and pick != "(upload below)":139 selected_from_outputs = Path(OUTPUTS) / pick140 uploaded = st.file_uploader("Or upload image for mesh", type=["png", "jpg"], key="mesh_upload") if use_image else None141 142 image_path_to_use = None143 if use_image and (selected_from_outputs and selected_from_outputs.exists() or uploaded):144 image_path_to_use = str(selected_from_outputs) if (selected_from_outputs and selected_from_outputs.exists()) else "upload"145 146 if st.button("Generate mesh", key="btn_mesh", disabled=not can_use_hub):147 if not image_path_to_use and not prompt_mesh.strip():148 st.warning("Enter a prompt or choose/upload an image.")149 else:150 with st.spinner("Running pipeline..."):151 try:152 from scripts.mesh_generator import (153 generate_mesh_from_image,154 generate_mesh_from_text,155 find_triposr_root,156 )157 triposr_root = find_triposr_root(str(ROOT))158 if not triposr_root:159 st.error(160 "TripoSR not found. In this Space the Docker image should include it. "161 "If you see this, rebuild the Space or check the Dockerfile."162 )163 elif image_path_to_use:164 path = image_path_to_use165 if path == "upload" and uploaded:166 path = os.path.join(OUTPUTS, "uploaded_mesh_input.png")167 with open(path, "wb") as f:168 f.write(uploaded.getvalue())169 if path != "upload" and os.path.isfile(path):170 import torch as _torch171 _dev = "cuda:0" if _torch.cuda.is_available() else "cpu"172 mesh_path, elapsed, msg = generate_mesh_from_image(173 path,174 output_dir=str(OUTPUTS / "mesh_run"),175 mesh_format=mesh_format,176 mc_resolution=mc_resolution,177 bake_texture=bake_texture,178 smooth_mesh=smooth_mesh,179 device=_dev,180 )181 if mesh_path:182 st.success(f"Done in {elapsed:.1f}s. {msg}")183 with open(mesh_path, "rb") as f:184 mesh_data = f.read()185 st.download_button("Download mesh", data=mesh_data, file_name=os.path.basename(mesh_path), key="dl_mesh_upload")186 else:187 st.error(msg)188 elif path == "upload":189 st.warning("Upload an image first.")190 else:191 import torch as _torch192 _dev = "cuda:0" if _torch.cuda.is_available() else "cpu"193 mesh_path, elapsed, msg = generate_mesh_from_text(194 prompt_mesh,195 output_dir=str(OUTPUTS),196 mesh_format=mesh_format,197 seed=seed_mesh,198 mc_resolution=mc_resolution,199 bake_texture=bake_texture,200 smooth_mesh=smooth_mesh,201 device=_dev,202 )203 if mesh_path:204 st.success(f"Done in {elapsed:.1f}s. {msg}")205 with open(mesh_path, "rb") as f:206 mesh_data = f.read()207 st.download_button("Download mesh", data=mesh_data, file_name=os.path.basename(mesh_path), key="dl_mesh")208 else:209 st.error(msg)210 except Exception as e:211 st.exception(e)212 213 # View 3D mesh (GLB): path, upload, or pick from outputs214 with st.expander("View 3D mesh", expanded=False):215 st.caption("Open a .glb file by path, upload, or pick from outputs. Drag to rotate, scroll to zoom.")216 from scripts.mesh_viewer import mesh_viewer_html217 import streamlit.components.v1 as components218 219 viewer_glb_path: str | None = None220 viewer_glb_bytes: bytes | None = None221 222 path_input = st.text_input(223 "Path to .glb file",224 value="",225 key="mesh_viewer_path",226 placeholder=r"e.g. C:\Users\...\Downloads\mesh (1).glb",227 )228 if path_input and Path(path_input.strip()).is_file():229 viewer_glb_path = path_input.strip()230 231 uploaded_glb = st.file_uploader("Or upload a .glb file", type=["glb"], key="mesh_viewer_upload")232 if uploaded_glb is not None:233 viewer_glb_bytes = uploaded_glb.getvalue()234 235 output_glbs = sorted(Path(OUTPUTS).rglob("*.glb"), key=lambda p: p.stat().st_mtime, reverse=True)236 if not viewer_glb_path and not viewer_glb_bytes and output_glbs:237 opt_names = [str(p.relative_to(OUTPUTS)) for p in output_glbs]238 k = "mesh_viewer_pick"239 if k in st.session_state and st.session_state[k] not in opt_names:240 del st.session_state[k]241 picked = st.selectbox("Or pick from outputs", ["(none)"] + opt_names, key=k)242 if picked and picked != "(none)":243 viewer_glb_path = str(OUTPUTS / picked)244 245 if viewer_glb_path or viewer_glb_bytes:246 html = mesh_viewer_html(glb_path=viewer_glb_path, glb_bytes=viewer_glb_bytes, height_px=480)247 components.html(html, height=500, scrolling=False)248 else:249 st.info("Enter a path to a .glb file, upload one, or generate a mesh above and pick it from outputs.")250 251 with tab_skybox:252 st.subheader("Generate 2:1 equirectangular skybox")253 st.markdown(254 "Uses **Stable Diffusion 2.1** at 2:1 aspect (e.g. 1024×512). "255 "Optional seamless check compares left/right edges."256 )257 prompt_sky = st.text_input(258 "Prompt (e.g. for skybox)",259 value="Cyberpunk city skyline at dusk, neon reflections, cinematic lighting.",260 key="sky_prompt",261 )262 col1, col2 = st.columns(2)263 with col1:264 width = st.selectbox("Width", [1024, 2048], key="sky_w")265 height = width // 2266 seed_sky = st.number_input("Seed (optional)", value=42, min_value=0, key="sky_seed")267 with col2:268 check_seamless = st.checkbox("Run seamless edge check", value=True, key="seamless")269 270 if st.button("Generate skybox", key="btn_sky", disabled=not can_use_hub):271 if not prompt_sky.strip():272 st.warning("Enter a prompt.")273 else:274 try:275 from scripts.skybox_generator import generate_skybox276 from scripts.check_seamless import check_seamless as run_seamless277 278 progress_placeholder = st.empty()279 status_placeholder = st.empty()280 progress_placeholder.progress(0)281 status_placeholder.caption("Loading model and starting generation…")282 283 def on_step(step: int, total: int) -> None:284 progress = min(step / total, 1.0)285 progress_placeholder.progress(progress)286 status_placeholder.caption(f"Step {min(step, total)} / {total}")287 288 out_path, elapsed, vram_mb = generate_skybox(289 prompt_sky,290 output_dir=str(OUTPUTS),291 width=width,292 height=height,293 seed=seed_sky,294 progress_callback=on_step,295 )296 progress_placeholder.progress(1.0)297 status_placeholder.caption("Done.")298 299 st.success(f"Done in {elapsed:.1f}s. Peak VRAM: {vram_mb:.0f} MB")300 st.image(out_path, use_container_width=True)301 with open(out_path, "rb") as f:302 skybox_data = f.read()303 st.download_button("Download skybox", data=skybox_data, file_name=os.path.basename(out_path), key="dl_sky")304 305 if check_seamless:306 result = run_seamless(out_path)307 st.info(result["message"])308 309 st.session_state["last_skybox_path"] = str(Path(out_path).resolve())310 except Exception as e:311 st.exception(e)312 313 # Show 360° viewer for last generated skybox (same session)314 if "last_skybox_path" in st.session_state:315 last_path = Path(st.session_state["last_skybox_path"]).resolve()316 if last_path.exists():317 with st.expander("View in 360°", expanded=False):318 st.caption("Drag to look around, scroll to zoom. Fullscreen available in the viewer.")319 from scripts.panorama_viewer import panorama_html320 import streamlit.components.v1 as components321 components.html(panorama_html(last_path, height_px=480), height=500, scrolling=False)322 323 # View existing image from outputs (or upload) in 360° – test without regenerating324 with st.expander("View existing image in 360°", expanded=False):325 st.caption("Pick an image from outputs or upload a 2:1 equirectangular image to test the viewer.")326 from scripts.panorama_viewer import panorama_html327 import streamlit.components.v1 as components328 329 output_files = sorted(Path(OUTPUTS).glob("*.png"), key=lambda p: p.stat().st_mtime, reverse=True)330 viewer_path = None331 option_names = [f.name for f in output_files]332 333 if output_files:334 key = "skybox_select_existing"335 if key in st.session_state and st.session_state[key] not in option_names:336 del st.session_state[key]337 selected_name = st.selectbox(338 "Choose image from outputs",339 options=option_names,340 key=key,341 )342 if selected_name:343 viewer_path = Path(OUTPUTS) / selected_name344 345 uploaded = st.file_uploader("Or upload a 2:1 equirectangular image", type=["png", "jpg", "jpeg"], key="skybox_upload_360")346 if uploaded is not None:347 upload_path = OUTPUTS / "uploaded_360_view.png"348 upload_path.write_bytes(uploaded.getvalue())349 viewer_path = upload_path350 351 if viewer_path is not None and viewer_path.exists():352 components.html(panorama_html(Path(viewer_path).resolve(), height_px=480), height=500, scrolling=False)353 elif not output_files and uploaded is None:354 st.info("No skybox images in outputs yet. Generate one above or upload an image.")355 356 st.divider()357 st.caption("Evoneural AI – Local ML Deployment MVP. Models run locally (no API).")358 359 360if __name__ == "__main__":361 main()362