Fix dia2 auto-install: clone from GitHub + editable pip install

dia2 is not on PyPI and has a packaging bug (missing subpackages).
Auto-installer now clones the repo and uses pip install -e --no-deps.
Also added version-agnostic TemplateResponse wrapper for Starlette
0.x/1.x compatibility (dia2 deps can downgrade starlette).

Author: devnen

.gitignore

@@ -36,6 +36,7 @@ env/
 outputs/
 reference_audio/
 model_cache/ # Also good practice to ignore the model cache
+dia2_src/ # Auto-cloned dia2 source (installed via editable pip)
 
 # Ignore test reports/coverage
 .coverage

engine.py

@@ -382,19 +382,51 @@ def get_compute_dtype(device: torch.device, weights_filename: str) -> str:
 
 
 def _auto_install_dia2():
-    """Attempts to install the dia2 package via pip at runtime."""
+    """
+    Attempts to install the dia2 package from GitHub at runtime.
+    dia2 is NOT on PyPI — must be cloned and installed in editable mode
+    because the pyproject.toml has a packaging bug (missing subpackages).
+    """
     global DIA2_AVAILABLE, Dia2, GenerationConfig, SamplingConfig, PrefixConfig, GenerationResult
     import subprocess
     import sys
 
+    DIA2_REPO_URL = "https://github.com/nari-labs/dia2.git"
+    dia2_src_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dia2_src")
+
     try:
         _check_cancelled()
+
+        # Step 1: Clone the dia2 repo if not already present
+        if not os.path.isdir(dia2_src_dir):
+            logger.info(f"Cloning dia2 repository from {DIA2_REPO_URL}...")
+            _update_download_status("downloading", "Cloning dia2 repository from GitHub...", 10)
+            result = subprocess.run(
+                ["git", "clone", DIA2_REPO_URL, dia2_src_dir],
+                capture_output=True, text=True, timeout=300,
+            )
+            if result.returncode != 0:
+                raise RuntimeError(f"git clone failed:\n{result.stderr}")
+            logger.info("dia2 repository cloned successfully.")
+        else:
+            logger.info(f"dia2 source already exists at {dia2_src_dir}, pulling latest...")
+            _update_download_status("downloading", "Updating dia2 repository...", 10)
+            subprocess.run(
+                ["git", "-C", dia2_src_dir, "pull", "--ff-only"],
+                capture_output=True, text=True, timeout=60,
+            )
+
+        _check_cancelled()
+
+        # Step 2: Install in editable mode with --no-deps to avoid breaking other packages
+        logger.info("Installing dia2 package (editable mode, no-deps)...")
+        _update_download_status("installing", "Installing dia2 package...", 25)
         result = subprocess.run(
-            [sys.executable, "-m", "pip", "install", "dia2"],
-            capture_output=True, text=True, timeout=600,
+            [sys.executable, "-m", "pip", "install", "-e", dia2_src_dir, "--no-deps"],
+            capture_output=True, text=True, timeout=300,
         )
         if result.returncode != 0:
-            raise RuntimeError(f"pip install dia2 failed:\n{result.stderr}")
+            raise RuntimeError(f"pip install -e dia2 failed:\n{result.stderr}")
 
         logger.info("dia2 package installed successfully. Importing...")
         from dia2 import Dia2 as _Dia2, GenerationConfig as _GC, SamplingConfig as _SC, PrefixConfig as _PC, GenerationResult as _GR
@@ -409,7 +441,7 @@ def _auto_install_dia2():
         logger.error(f"Failed to auto-install dia2: {e}", exc_info=True)
         raise ImportError(
             f"dia2 package could not be installed automatically: {e}. "
-            "Please install it manually: pip install dia2"
+            "Please install manually: git clone https://github.com/nari-labs/dia2.git dia2_src && pip install -e dia2_src --no-deps"
         )
 
 

server.py

@@ -221,6 +221,20 @@ def _delayed_browser_open(host, port):
 templates = Jinja2Templates(directory="ui")
 
 
+def _render_template(request, name, context, status_code=200):
+    """Version-agnostic TemplateResponse wrapper for Starlette 0.x and 1.x."""
+    import starlette
+    major = int(starlette.__version__.split(".")[0])
+    if major >= 1:
+        # Starlette 1.0+: request is a separate positional arg
+        return templates.TemplateResponse(request, name, context=context, status_code=status_code)
+    else:
+        # Starlette 0.x: request goes inside the context dict
+        ctx = dict(context) if context else {}
+        ctx["request"] = request
+        return templates.TemplateResponse(name, ctx, status_code=status_code)
+
+
 # --- Configuration Routes (New YAML-based) ---
 @app.post(
     "/save_settings",
@@ -831,7 +845,7 @@ async def get_web_ui(request: Request):
             "success": None,
         }
 
-        return templates.TemplateResponse(request, "index.html", context=template_context)
+        return _render_template(request, "index.html", template_context)
 
     except Exception as e:
         logger.error(f"Error rendering Web UI: {e}", exc_info=True)
@@ -909,7 +923,7 @@ async def handle_web_ui_generate(
             "output_file_url": None,
             "generation_time": None,
         }
-        return templates.TemplateResponse(request, "index.html", context=error_context, status_code=503)
+        return _render_template(request, "index.html", error_context, status_code=503)
 
     # --- Start processing the valid request ---
     logger.info(
@@ -1052,9 +1066,7 @@ async def handle_web_ui_generate(
             "output_file_url": None,
             "generation_time": None,
         }
-        return templates.TemplateResponse(
-            request, "index.html", context=error_context, status_code=400
-        )  # Bad Request
+        return _render_template(request, "index.html", error_context, status_code=400)
 
     # --- Generation ---
     try:
@@ -1218,12 +1230,7 @@ async def handle_web_ui_generate(
     }
 
     # Render and return the HTML response
-    return templates.TemplateResponse(
-        request,
-        "index.html",
-        context=template_context,
-        status_code=status_code,
-    )
+    return _render_template(request, "index.html", template_context, status_code=status_code)
 
 
 # --- Reference Audio Upload Endpoint ---