Upload 2 files

app.py

@@ -5,36 +5,39 @@ import threading
 from pathlib import Path
 
 # === Runtime environment tweaks ===
-os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
-os.environ.setdefault("WORLDGEN_DISABLE_EVAL", "1")
-os.environ.setdefault("GRADIO_TEMP_DIR", "/tmp/gradio")
+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 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.utils
+
+    def safe_get_precision(*args, **kwargs):
+        return "fp32"  # force CPU precision
+
+    nunchaku.utils.get_precision = safe_get_precision
+
     import nunchaku.models.transformers.transformer_flux as tf
     tf.get_precision = safe_get_precision
 except Exception:
     pass
 
 import gradio as gr
-from huggingface_hub import hf_hub_download
 from worldgen import WorldGen
 
+# Lazy init
 _wg = None
 def get_worldgen():
     global _wg
     if _wg is None:
-        # Confirm model file is downloadable
-        model_path = hf_hub_download(
-            repo_id="mit-han-lab/svdq-int4-flux.1-dev",
-            filename="unquantized_layers.safetensors",
-            token=None
+        _wg = WorldGen(
+            mode="t2s",
+            device="cpu",
+            low_vram=True,
+            lora_path=None,
+            repo_id="mit-han-lab/svdq-int4-flux.1-dev"
         )
-        _wg = WorldGen(mode="t2s", device="cpu", low_vram=True, lora_path=None)
     return _wg
 
 def _zip_directory(dir_path: Path) -> str:
@@ -46,15 +49,18 @@ def _zip_directory(dir_path: Path) -> str:
 
 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."
     )
@@ -63,20 +69,36 @@ def _save_scene(generated, out_dir: Path) -> Path:
 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)
-    threading.Timer(300, shutil.rmtree, args=(tmpdir,)).start()
+
+    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")
+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(
@@ -88,5 +110,7 @@ with gr.Blocks(title="WorldGen • Text -> 3D Scene") as demo:
         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()