Upload 2 files

app.py

@@ -4,12 +4,27 @@ import tempfile
 import threading
 from pathlib import Path
 
-# === Runtime environment tweaks ===
+# === Runtime environment tweaks (must be set before importing heavy libs) ===
 os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")  # force CPU
-os.environ.setdefault("WORLDGEN_DISABLE_EVAL", "1")
-os.environ.setdefault("GRADIO_TEMP_DIR", "/tmp/gradio")
+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"  # force CPU precision
+
+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
+import torch
 from worldgen import WorldGen
 
 _wg = None
@@ -28,15 +43,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."
     )
@@ -45,31 +63,48 @@ 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)
+
     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\n- CPU-only fallback is enabled.\n- Outputs are zipped for download.")
+    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", lines=3)
+            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()