Upload 2 files

app.py

@@ -4,16 +4,16 @@ import tempfile
 import threading
 from pathlib import Path
 
-# === Runtime environment tweaks (must be set before importing heavy libs) ===
+# === 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
+os.environ.setdefault("WORLDGEN_DISABLE_EVAL", "1")
+os.environ.setdefault("GRADIO_TEMP_DIR", "/tmp/gradio")
 
 # Monkey patch to avoid CUDA errors
 import nunchaku.utils
 
 def safe_get_precision(*args, **kwargs):
-    return "fp32"  # Force CPU precision
+    return "fp32"
 
 nunchaku.utils.get_precision = safe_get_precision
 
@@ -24,13 +24,19 @@ except Exception:
     pass
 
 import gradio as gr
-import torch
+from huggingface_hub import hf_hub_download
 from worldgen import WorldGen
 
 _wg = None
 def get_worldgen():
     global _wg
     if _wg is None:
+        # Download model file from correct repo
+        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)
     return _wg
 
@@ -43,18 +49,15 @@ 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,48 +66,27 @@ 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()
-
+    threading.Timer(300, shutil.rmtree, 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.
-"""
-    )
+    gr.Markdown("# WorldGen - Text -> 3D\n- CPU-only fallback is enabled.\n- Outputs are zipped for download.")
     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,
-            )
+            prompt = gr.Textbox(label="Scene prompt (English)", placeholder="a neon-lit cyberpunk street", 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()