worldgenTest / app.py
chahui's picture
Upload 2 files
9f4131d verified
import os
import shutil
import tempfile
import threading
from pathlib import Path
# === Runtime environment tweaks ===
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") # force CPU
os.environ.setdefault("WORLDGEN_DISABLE_EVAL", "1") # avoid KNN/eval paths if supported
os.environ.setdefault("GRADIO_TEMP_DIR", "/tmp/gradio") # centralize gradio cache
# Monkey patch get_precision to avoid CUDA errors
import nunchaku.utils
def safe_get_precision(*args, **kwargs):
return "fp32"
nunchaku.utils.get_precision = safe_get_precision
try:
import nunchaku.models.transformers.transformer_flux as tf
tf.get_precision = safe_get_precision
except Exception:
pass
import gradio as gr
from worldgen import WorldGen
_wg = None
def get_worldgen():
global _wg
if _wg is None:
_wg = WorldGen(mode="t2s", device="cpu", low_vram=True, lora_path=None)
return _wg
def _zip_directory(dir_path: Path) -> str:
zip_path = dir_path.with_suffix(".zip")
if zip_path.exists():
zip_path.unlink()
shutil.make_archive(str(dir_path), "zip", root_dir=str(dir_path))
return str(zip_path)
def _save_scene(generated, out_dir: Path) -> Path:
out_dir.mkdir(parents=True, exist_ok=True)
for meth in ("save", "export", "write", "to_disk"):
if hasattr(generated, meth):
getattr(generated, meth)(str(out_dir))
return out_dir
for meth in ("save_as_glb", "to_glb", "save_as_ply", "to_ply"):
if hasattr(generated, meth):
fp = out_dir / ("scene.glb" if "glb" in meth else "scene.ply")
getattr(generated, meth)(str(fp))
return out_dir
(out_dir / "README.txt").write_text(
"No known save/export method detected. Please adjust _save_scene according to your WorldGen version."
)
return out_dir
def generate_from_text(prompt: str):
if not prompt or not prompt.strip():
return None, "Please enter a scene description (English works best)."
wg = get_worldgen()
try:
result = wg.generate_world(prompt.strip())
except Exception as e:
return None, f"Generation failed: {e}"
tmpdir = Path(tempfile.mkdtemp())
out_dir = tmpdir / "worldgen_output"
out_dir = _save_scene(result, out_dir)
zip_path = _zip_directory(out_dir)
def _cleanup(p: Path):
try:
shutil.rmtree(p, ignore_errors=True)
except Exception:
pass
threading.Timer(300, _cleanup, args=(tmpdir,)).start()
return zip_path, "Done. Download the ZIP (contains 3D assets)."
with gr.Blocks(title="WorldGen • Text -> 3D Scene") as demo:
gr.Markdown(
"""# WorldGen - Text -> 3D
- CPU-only fallback is enabled. For real runs, switch the Space to a GPU (T4/A10G).
- Outputs are zipped for download; temp files are cleaned shortly after each run.
"""
)
with gr.Row():
with gr.Column():
prompt = gr.Textbox(
label="Scene prompt (English)",
placeholder="a neon-lit cyberpunk street with rain reflections, highly detailed",
lines=3,
)
btn = gr.Button("Generate 3D Scene")
with gr.Column():
out_zip = gr.File(label="Download output (ZIP)")
status = gr.Markdown()
btn.click(generate_from_text, inputs=[prompt], outputs=[out_zip, status])
demo.launch()