Enhance app.py: Add support for MuseScore 3 binary, improve rendering process, and persist rendered images

app.py

@@ -49,11 +49,22 @@ CPU_COUNT = os.cpu_count() or 1
 MUSESCORE_APPIMAGE_URL = "https://www.modelscope.cn/studio/Genius-Society/piano_trans/resolve/master/MuseScore.AppImage"
 MUSESCORE_STORAGE_DIR = Path("artifacts") / "musescore"
 MUSESCORE_ENV_VAR = "MUSESCORE_BIN"
-MUSESCORE_RENDER_TIMEOUT = int(os.environ.get("MUSESCORE_RENDER_TIMEOUT", "120"))
+MUSESCORE_RENDER_TIMEOUT = int(os.environ.get("MUSESCORE_RENDER_TIMEOUT", "180"))
 MUSESCORE_EXTRACT_TIMEOUT = int(os.environ.get("MUSESCORE_EXTRACT_TIMEOUT", "240"))
 _MUSESCORE_BINARY: Optional[str] = None
 _MUSESCORE_READY: bool = False
 
+MUSESCORE_V3_APPIMAGE_URL = "https://github.com/musescore/MuseScore/releases/download/v3.6.2/MuseScore-3.6.2.548021370-x86_64.AppImage"
+MUSESCORE_V3_STORAGE_DIR = Path("artifacts") / "musescore_v3"
+MUSESCORE_V3_ENV_VAR = "MUSESCORE_V3_BIN"
+_MUSESCORE_V3_BINARY: Optional[str] = None
+
+RENDER_OUTPUT_DIR = Path("artifacts") / "rendered_scores"
+
+XVFB_ENV_VAR = "XVFB_BIN"
+XVFB_STORAGE_DIR = Path("artifacts") / "xvfb"
+_XVFB_BINARY: Optional[str] = None
+
 # Global model variable
 MODEL = None
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -295,17 +306,172 @@ def ensure_musescore_binary() -> Optional[str]:
     return None
 
 
+def ensure_musescore_v3_binary() -> Optional[str]:
+    """Ensure a MuseScore 3.x binary is available for rendering."""
+    global _MUSESCORE_V3_BINARY
+    if _MUSESCORE_V3_BINARY and os.path.exists(_MUSESCORE_V3_BINARY):
+        return _MUSESCORE_V3_BINARY
+    env_path = os.environ.get(MUSESCORE_V3_ENV_VAR)
+    if env_path and os.path.exists(env_path):
+        logger.info("Using MuseScore 3 binary from %s", MUSESCORE_V3_ENV_VAR)
+        _MUSESCORE_V3_BINARY = env_path
+        return _MUSESCORE_V3_BINARY
+    storage = MUSESCORE_V3_STORAGE_DIR
+    storage.mkdir(parents=True, exist_ok=True)
+    appimage_path = (storage / "MuseScore-3.AppImage").resolve(strict=False)
+    apprun_path = (storage / "squashfs-root" / "AppRun").resolve(strict=False)
+    if apprun_path.exists():
+        logger.info("Using cached MuseScore 3 AppRun at %s", apprun_path)
+        _MUSESCORE_V3_BINARY = str(apprun_path)
+        return _MUSESCORE_V3_BINARY
+    if not appimage_path.exists():
+        logger.info("MuseScore 3 AppImage missing. Downloading (first run only)...")
+        if not _download_file(MUSESCORE_V3_APPIMAGE_URL, appimage_path):
+            return None
+    try:
+        os.chmod(appimage_path, 0o755)
+    except Exception as exc:
+        logger.warning("Could not chmod MuseScore 3 AppImage: %s", exc)
+    try:
+        with log_timing("Extracting MuseScore 3 AppImage"):
+            subprocess.run(
+                [str(appimage_path), "--appimage-extract"],
+                cwd=storage,
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                timeout=MUSESCORE_EXTRACT_TIMEOUT,
+            )
+    except subprocess.CalledProcessError as exc:
+        stderr = exc.stderr.decode(errors='ignore') if exc.stderr else str(exc)
+        logger.error("MuseScore 3 extraction failed: %s", stderr)
+        return None
+    except subprocess.TimeoutExpired:
+        logger.error("MuseScore 3 extraction timed out after %ss", MUSESCORE_EXTRACT_TIMEOUT)
+        return None
+    if apprun_path.exists():
+        _MUSESCORE_V3_BINARY = str(apprun_path)
+        try:
+            os.chmod(apprun_path, 0o755)
+        except Exception:
+            pass
+        logger.info("MuseScore 3 AppRun ready at %s", _MUSESCORE_V3_BINARY)
+        return _MUSESCORE_V3_BINARY
+    logger.error("MuseScore 3 extraction did not produce the expected AppRun binary.")
+    return None
+
+
+def _download_xvfb_package(dest_dir: Path) -> Optional[Path]:
+    """Download the Xvfb .deb package using apt."""
+    try:
+        completed = subprocess.run(
+            ["apt", "download", "xvfb"],
+            cwd=str(dest_dir),
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        logger.debug("apt download xvfb stdout: %s", completed.stdout.strip())
+        if completed.stderr:
+            logger.debug("apt download xvfb stderr: %s", completed.stderr.strip())
+    except FileNotFoundError:
+        logger.error("'apt' command not available; cannot download Xvfb automatically.")
+        return None
+    except subprocess.CalledProcessError as exc:
+        logger.error(
+            "Failed to download Xvfb package (exit %s): %s",
+            exc.returncode,
+            exc.stderr.strip() if exc.stderr else exc,
+        )
+        return None
+    deb_candidates = sorted(dest_dir.glob("xvfb_*.deb"), key=lambda p: p.stat().st_mtime, reverse=True)
+    if not deb_candidates:
+        logger.error("apt download xvfb did not produce any .deb files under %s", dest_dir)
+        return None
+    return deb_candidates[0]
+
+
+def _extract_xvfb_binary(deb_path: Path, target_dir: Path) -> Optional[Path]:
+    extract_dir = target_dir / "pkg"
+    if extract_dir.exists():
+        shutil.rmtree(extract_dir, ignore_errors=True)
+    try:
+        subprocess.run(
+            ["dpkg-deb", "-x", str(deb_path), str(extract_dir)],
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+    except FileNotFoundError:
+        logger.error("'dpkg-deb' command not available; cannot extract Xvfb package.")
+        return None
+    except subprocess.CalledProcessError as exc:
+        stderr = exc.stderr.decode(errors="ignore") if isinstance(exc.stderr, bytes) else exc.stderr
+        logger.error("Failed to extract Xvfb package: %s", stderr or exc)
+        return None
+    xvfb_path = extract_dir / "usr/bin/Xvfb"
+    if xvfb_path.exists():
+        logger.info("Xvfb binary extracted to %s", xvfb_path)
+        try:
+            os.chmod(xvfb_path, 0o755)
+        except Exception:
+            pass
+        try:
+            deb_path.unlink()
+        except Exception:
+            pass
+        return xvfb_path
+    logger.error("Extracted Xvfb package but could not find usr/bin/Xvfb inside %s", extract_dir)
+    return None
+
+
+def ensure_xvfb_binary() -> Optional[str]:
+    """Ensure we have an Xvfb binary available for headless rendering."""
+    global _XVFB_BINARY
+    if _XVFB_BINARY and os.path.exists(_XVFB_BINARY):
+        return _XVFB_BINARY
+    env_path = os.environ.get(XVFB_ENV_VAR)
+    if env_path and os.path.exists(env_path):
+        _XVFB_BINARY = env_path
+        return _XVFB_BINARY
+    which = shutil.which("Xvfb")
+    if which:
+        _XVFB_BINARY = which
+        return _XVFB_BINARY
+    XVFB_STORAGE_DIR.mkdir(parents=True, exist_ok=True)
+    extracted_bin = XVFB_STORAGE_DIR / "pkg" / "usr" / "bin" / "Xvfb"
+    if extracted_bin.exists():
+        _XVFB_BINARY = str(extracted_bin)
+        return _XVFB_BINARY
+    deb_path = _download_xvfb_package(XVFB_STORAGE_DIR)
+    if not deb_path:
+        return None
+    extracted = _extract_xvfb_binary(deb_path, XVFB_STORAGE_DIR)
+    if extracted:
+        _XVFB_BINARY = str(extracted)
+        return _XVFB_BINARY
+    return None
+
+
 def initialize_musescore_backend() -> bool:
     """Initialize MuseScore AppRun at startup to avoid on-demand downloads."""
     global _MUSESCORE_READY
     if _MUSESCORE_READY:
         return True
-    binary = ensure_musescore_binary()
-    if binary:
-        logger.info("MuseScore AppRun ready at startup: %s", binary)
+    available = []
+    primary = ensure_musescore_binary()
+    if primary:
+        available.append(primary)
+        logger.info("MuseScore 4 AppRun ready at startup: %s", primary)
+    legacy = ensure_musescore_v3_binary()
+    if legacy:
+        available.append(legacy)
+        logger.info("MuseScore 3 AppRun ready at startup: %s", legacy)
+    if available:
         _MUSESCORE_READY = True
         return True
-    logger.warning("MuseScore AppRun could not be initialized; score visualization will fail.")
+    logger.warning("No MuseScore AppRun binaries could be initialized; score visualization will fail.")
     return False
 
 
@@ -355,42 +521,132 @@ def _coalesce_musescore_output(output_path: str) -> Optional[str]:
     return str(normalized_path)
 
 
+def persist_rendered_image(src_path: str) -> Optional[str]:
+    """Copy rendered PNG to a persistent artifacts directory for UI display."""
+    if not src_path or not os.path.exists(src_path):
+        return None
+    try:
+        RENDER_OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+        dest = RENDER_OUTPUT_DIR / f"{int(time.time()*1000)}_{Path(src_path).name}"
+        shutil.copy2(src_path, dest)
+        return str(dest)
+    except Exception as exc:
+        logger.warning("Could not persist rendered image %s: %s", src_path, exc)
+        return src_path
+
+
+@contextmanager
+def xvfb_session():
+    """Spin up a temporary Xvfb server if available."""
+    xvfb_bin = ensure_xvfb_binary()
+    if not xvfb_bin:
+        logger.warning("Xvfb binary unavailable; proceeding without virtual display.")
+        yield None
+        return
+    display = None
+    base_dir = Path("/tmp/.X11-unix")
+    try:
+        base_dir.mkdir(mode=0o1777, exist_ok=True)
+    except Exception:
+        pass
+    used = {p.name for p in base_dir.glob("X*")}
+    for candidate in range(99, 160):
+        name = f"X{candidate}"
+        if name not in used:
+            display = f":{candidate}"
+            break
+    if display is None:
+        logger.warning("No free DISPLAY slots for Xvfb.")
+        yield None
+        return
+    cmd = [
+        xvfb_bin,
+        display,
+        "-screen",
+        "0",
+        "1920x1080x24",
+        "-nolisten",
+        "tcp",
+    ]
+    logger.debug("Starting Xvfb with command: %s", " ".join(cmd))
+    proc = subprocess.Popen(
+        cmd,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    time.sleep(0.5)
+    if proc.poll() is not None:
+        logger.error("Xvfb failed to start (exit %s).", proc.returncode)
+        yield None
+        return
+    try:
+        yield display
+    finally:
+        proc.terminate()
+        try:
+            proc.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+
+
 def render_with_musescore(musicxml_path: Optional[str], output_path: str) -> Optional[str]:
     """Render using MuseScore command-line interface."""
     if not musicxml_path or not os.path.exists(musicxml_path):
         return None
-    musescore_bin = ensure_musescore_binary()
-    if musescore_bin is None:
-        logger.warning("MuseScore binary unavailable; skipping MuseScore rendering.")
-        return None
-    env = os.environ.copy()
-    env["QT_QPA_PLATFORM"] = "offscreen"
-    env.setdefault("QTWEBENGINE_DISABLE_SANDBOX", "1")
-    env.setdefault("MUSESCORE_NO_AUDIO", "1")
-    cmd = [musescore_bin, "-o", output_path, musicxml_path]
-    logger.debug("Executing MuseScore command: %s", " ".join(cmd))
-    try:
-        with log_timing("MuseScore rendering"):
-            subprocess.run(
-                cmd,
-                check=True,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                env=env,
-                timeout=MUSESCORE_RENDER_TIMEOUT,
-            )
-    except subprocess.CalledProcessError as exc:
-        stderr = exc.stderr.decode(errors='ignore') if exc.stderr else str(exc)
-        logger.error("MuseScore rendering failed: %s", stderr)
-        return None
-    except subprocess.TimeoutExpired:
-        logger.error("MuseScore rendering timed out after %ss", MUSESCORE_RENDER_TIMEOUT)
+    candidates = []
+    legacy = ensure_musescore_v3_binary()
+    if legacy:
+        candidates.append(("MuseScore 3", legacy, True))
+    primary = ensure_musescore_binary()
+    if primary:
+        candidates.append(("MuseScore 4", primary, True))
+    if not candidates:
+        logger.warning("No MuseScore binaries available for rendering.")
         return None
-    normalized_path = _coalesce_musescore_output(output_path)
-    if normalized_path and os.path.exists(normalized_path):
-        logger.info("MuseScore rendered %s -> %s", musicxml_path, normalized_path)
-        return normalized_path
-    logger.error("MuseScore rendered score but the expected output file was not found.")
+    last_error = None
+    for label, musescore_bin, requires_display in candidates:
+        env = os.environ.copy()
+        env.setdefault("QTWEBENGINE_DISABLE_SANDBOX", "1")
+        env.setdefault("MUSESCORE_NO_AUDIO", "1")
+        cmd = [musescore_bin, "-o", output_path, musicxml_path]
+        logger.info("Attempting rendering with %s (%s).", label, musescore_bin)
+        try:
+            with xvfb_session() as display:
+                if display:
+                    env["DISPLAY"] = display
+                    env["QT_QPA_PLATFORM"] = "xcb"
+                    logger.debug("%s: using Xvfb display %s", label, display)
+                else:
+                    if requires_display:
+                        logger.warning("%s requires an X11 display but Xvfb could not be started.", label)
+                        continue
+                    env["QT_QPA_PLATFORM"] = "offscreen"
+                    logger.debug("%s: using Qt offscreen platform.", label)
+                with log_timing(f"{label} rendering"):
+                    subprocess.run(
+                        cmd,
+                        check=True,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.PIPE,
+                        env=env,
+                        timeout=MUSESCORE_RENDER_TIMEOUT,
+                    )
+        except subprocess.CalledProcessError as exc:
+            stderr = exc.stderr.decode(errors='ignore') if exc.stderr else str(exc)
+            logger.error("%s rendering failed: %s", label, stderr)
+            last_error = stderr
+            continue
+        except subprocess.TimeoutExpired:
+            logger.error("%s rendering timed out after %ss", label, MUSESCORE_RENDER_TIMEOUT)
+            last_error = f"{label} timed out"
+            continue
+        normalized_path = _coalesce_musescore_output(output_path)
+        if normalized_path and os.path.exists(normalized_path):
+            logger.info("%s rendered %s -> %s", label, musicxml_path, normalized_path)
+            return normalized_path
+        logger.error("%s rendered score but the expected output file was not found.", label)
+        last_error = "output missing"
+    logger.error("All MuseScore binaries failed to render the score. Last error: %s", last_error)
     return None
 
 
@@ -630,8 +886,9 @@ def process_musicxml(
                 source_musicxml_path=source_path,
             )
             predictions = predict_analysis(model, score, selected_tasks)
-        if rendered_path is None:
-            logger.warning("MuseScore AppRun could not render the score. Visualization will be unavailable.")
+        persisted_path = persist_rendered_image(rendered_path) if rendered_path else None
+        if rendered_path is None or persisted_path is None:
+            logger.warning("MuseScore AppRun could not render the score or save the PNG; visualization will be unavailable.")
         
         # Create DataFrame
         if predictions:
@@ -684,7 +941,7 @@ def process_musicxml(
             if parsed_score_path:
                 status_msg += " Parsed MusicXML ready for download."
         
-        return rendered_path, df, parsed_score_path, status_msg
+        return persisted_path, df, parsed_score_path, status_msg
         
     except Exception as e:
         error_msg = f"Error processing file: {str(e)}\n\n{traceback.format_exc()}"