Add GPU ComfyUI diagnostic tests

app.py

@@ -1,40 +1,340 @@
+from __future__ import annotations
+
+import json
+import math
+import subprocess
+import sys
+import time
+import uuid
+import wave
+from pathlib import Path
+from typing import Any
+
 import gradio as gr
+import requests
 import spaces
 import torch
 
 
-@spaces.GPU(duration=60)
-def gpu_smoke_test():
-    lines = ["VoiceGate GPU smoke test"]
+ROOT = Path(__file__).resolve().parent
+COMFY_DIR = ROOT / "ComfyUI"
+COMFY_INPUT_DIR = COMFY_DIR / "input"
+COMFY_LOG = Path("/tmp/voicegate_comfy_gradio.log")
+COMFY_URL = "http://127.0.0.1:8188"
+COMFY_HOST = "127.0.0.1"
+COMFY_PORT = "8188"
+
+COMFY_PROCESS: subprocess.Popen | None = None
+BOOTSTRAPPED = False
+
+
+def gpu_status_lines() -> list[str]:
+    lines = ["VoiceGate GPU status"]
     lines.append(f"torch={torch.__version__}")
     lines.append(f"cuda_available={torch.cuda.is_available()}")
     lines.append(f"cuda_device_count={torch.cuda.device_count()}")
+    if torch.cuda.is_available():
+        props = torch.cuda.get_device_properties(0)
+        lines.append(f"device_name={torch.cuda.get_device_name(0)}")
+        lines.append(f"total_memory_gb={props.total_memory / 1024**3:.2f}")
+    return lines
+
+
+def wait_for_comfy(timeout: float = 180) -> dict[str, Any]:
+    deadline = time.time() + timeout
+    last_error = ""
+    while time.time() < deadline:
+        try:
+            response = requests.get(f"{COMFY_URL}/system_stats", timeout=5)
+            if response.ok:
+                return response.json()
+            last_error = f"HTTP {response.status_code}: {response.text[:300]}"
+        except requests.RequestException as exc:
+            last_error = repr(exc)
+        time.sleep(2)
+    raise RuntimeError(f"ComfyUI did not become ready: {last_error}")
+
+
+def run_bootstrap(lines: list[str]) -> None:
+    global BOOTSTRAPPED
+
+    if BOOTSTRAPPED and COMFY_DIR.exists():
+        lines.append("bootstrap=already_done")
+        return
+
+    started = time.time()
+    lines.append("bootstrap=starting")
+    command = [sys.executable, str(ROOT / "scripts" / "bootstrap_comfy.py")]
+    result = subprocess.run(
+        command,
+        cwd=ROOT,
+        text=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        timeout=900,
+    )
+    lines.append(f"bootstrap_returncode={result.returncode}")
+    lines.append(f"bootstrap_elapsed_sec={time.time() - started:.1f}")
+    if result.returncode != 0:
+        lines.append("bootstrap_tail:")
+        lines.extend(result.stdout.splitlines()[-80:])
+        raise RuntimeError("bootstrap_comfy.py failed")
+    BOOTSTRAPPED = True
+
+
+def ensure_comfy(lines: list[str], *, timeout: float = 240) -> dict[str, Any]:
+    global COMFY_PROCESS
+
+    run_bootstrap(lines)
+
+    try:
+        stats = wait_for_comfy(timeout=5)
+        lines.append("comfy=already_running")
+        return stats
+    except RuntimeError:
+        pass
+
+    log = COMFY_LOG.open("ab")
+    command = [
+        sys.executable,
+        "main.py",
+        "--listen",
+        COMFY_HOST,
+        "--port",
+        COMFY_PORT,
+    ]
+    COMFY_PROCESS = subprocess.Popen(
+        command,
+        cwd=COMFY_DIR,
+        stdout=log,
+        stderr=subprocess.STDOUT,
+    )
+    lines.append(f"comfy_started_pid={COMFY_PROCESS.pid}")
+    try:
+        return wait_for_comfy(timeout=timeout)
+    except Exception:
+        lines.append("comfy_log_tail:")
+        if COMFY_LOG.exists():
+            lines.extend(COMFY_LOG.read_text(encoding="utf-8", errors="replace").splitlines()[-120:])
+        raise
+
+
+def write_sine_wav(filename: str, *, seconds: float = 1.0, frequency: float = 440.0) -> str:
+    COMFY_INPUT_DIR.mkdir(parents=True, exist_ok=True)
+    path = COMFY_INPUT_DIR / filename
+    sample_rate = 16000
+    total = int(sample_rate * seconds)
+    amplitude = 0.2
+    with wave.open(str(path), "wb") as file:
+        file.setnchannels(1)
+        file.setsampwidth(2)
+        file.setframerate(sample_rate)
+        for index in range(total):
+            value = int(32767 * amplitude * math.sin(2 * math.pi * frequency * index / sample_rate))
+            file.writeframesraw(value.to_bytes(2, byteorder="little", signed=True))
+    return filename
+
+
+def submit_prompt(workflow: dict[str, Any]) -> str:
+    response = requests.post(
+        f"{COMFY_URL}/prompt",
+        json={"prompt": workflow, "client_id": str(uuid.uuid4())},
+        timeout=120,
+    )
+    response.raise_for_status()
+    return response.json()["prompt_id"]
+
+
+def wait_for_history(prompt_id: str, timeout: float = 1200) -> dict[str, Any]:
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        response = requests.get(f"{COMFY_URL}/history/{prompt_id}", timeout=30)
+        response.raise_for_status()
+        payload = response.json()
+        if prompt_id in payload:
+            return payload[prompt_id]
+        time.sleep(2)
+    raise TimeoutError(f"Timed out waiting for prompt {prompt_id}")
+
+
+def history_summary(history: dict[str, Any]) -> list[str]:
+    lines = []
+    status = history.get("status", {})
+    lines.append(f"status_str={status.get('status_str')}")
+    lines.append(f"completed={status.get('completed')}")
+    messages = status.get("messages") or []
+    errors = [message for message in messages if isinstance(message, list) and message[0] == "execution_error"]
+    if errors:
+        lines.append("errors:")
+        lines.append(json.dumps(errors, ensure_ascii=False, indent=2)[:4000])
+
+    outputs = history.get("outputs", {})
+    output_files = []
+    for node_output in outputs.values():
+        for key in ("audio", "images", "gifs"):
+            for item in node_output.get(key, []) or []:
+                filename = item.get("filename")
+                subfolder = item.get("subfolder")
+                if subfolder:
+                    output_files.append(f"{subfolder}/{filename}")
+                elif filename:
+                    output_files.append(filename)
+    if output_files:
+        lines.append("outputs:")
+        lines.extend(output_files)
+    return lines
+
+
+def melband_workflow(audio_filename: str, prefix: str) -> dict[str, Any]:
+    return {
+        "1": {
+            "class_type": "LoadAudio",
+            "inputs": {"audio": audio_filename, "audioUI": ""},
+        },
+        "2": {
+            "class_type": "MelBandRoFormerModelLoader",
+            "inputs": {"model_name": "MelBandRoFormer_comfy/MelBandRoformer_fp32.safetensors"},
+        },
+        "3": {
+            "class_type": "MelBandRoFormerSampler",
+            "inputs": {"model": ["2", 0], "audio": ["1", 0]},
+        },
+        "4": {
+            "class_type": "SaveAudioMP3",
+            "inputs": {
+                "filename_prefix": f"audio/{prefix}_vocals",
+                "quality": "V0",
+                "audioUI": "",
+                "audio": ["3", 0],
+            },
+        },
+        "5": {
+            "class_type": "SaveAudioMP3",
+            "inputs": {
+                "filename_prefix": f"audio/{prefix}_instruments",
+                "quality": "V0",
+                "audioUI": "",
+                "audio": ["3", 1],
+            },
+        },
+    }
+
+
+def voxcpm_tts_workflow(prefix: str) -> dict[str, Any]:
+    return {
+        "1": {
+            "class_type": "RunningHub_VoxCPM_LoadModel",
+            "inputs": {"model_name": "VoxCPM2", "optimize": False, "lora_name": "None"},
+        },
+        "2": {
+            "class_type": "RunningHub_VoxCPM_Generate",
+            "inputs": {
+                "model": ["1", 0],
+                "control_instruction": "清晰自然的中文女声",
+                "text": "你好，VoiceGate GPU 语音合成测试。",
+                "cfg_value": 2.0,
+                "inference_steps": 4,
+                "seed": 20260605,
+                "ultimate_clone": False,
+                "reference_audio_text": "",
+                "normalize_text": False,
+                "denoise_reference": False,
+                "max_len": 512,
+                "retry_badcase": True,
+            },
+        },
+        "3": {
+            "class_type": "SaveAudioMP3",
+            "inputs": {
+                "filename_prefix": f"audio/{prefix}",
+                "quality": "V0",
+                "audioUI": "",
+                "audio": ["2", 0],
+            },
+        },
+    }
+
+
+@spaces.GPU(duration=60)
+def gpu_smoke_test() -> str:
+    lines = gpu_status_lines()
+    if torch.cuda.is_available():
+        tensor = torch.arange(16, device="cuda:0", dtype=torch.float32)
+        result = (tensor * 2).sum().item()
+        torch.cuda.synchronize()
+        lines.append(f"tensor_result={result}")
+        lines.append(f"memory_reserved_mb={torch.cuda.memory_reserved(0) / 1024**2:.2f}")
+    return "\n".join(lines)
+
+
+@spaces.GPU(duration=900)
+def comfy_runtime_test() -> str:
+    lines = gpu_status_lines()
+    started = time.time()
+    try:
+        stats = ensure_comfy(lines)
+        lines.append(f"comfy_ready=true")
+        lines.append(f"comfy_elapsed_sec={time.time() - started:.1f}")
+        lines.append("system_stats:")
+        lines.append(json.dumps(stats, ensure_ascii=False, indent=2)[:4000])
+    except Exception as exc:
+        lines.append(f"error={type(exc).__name__}: {exc}")
+    return "\n".join(lines)
+
+
+@spaces.GPU(duration=1200)
+def melband_gpu_test() -> str:
+    lines = gpu_status_lines()
+    started = time.time()
+    try:
+        ensure_comfy(lines)
+        audio_filename = write_sine_wav(f"voicegate_melband_{uuid.uuid4().hex[:8]}.wav")
+        prefix = f"melband_gpu_{uuid.uuid4().hex[:8]}"
+        prompt_id = submit_prompt(melband_workflow(audio_filename, prefix))
+        lines.append(f"prompt_id={prompt_id}")
+        history = wait_for_history(prompt_id)
+        lines.extend(history_summary(history))
+        lines.append(f"elapsed_sec={time.time() - started:.1f}")
+    except Exception as exc:
+        lines.append(f"error={type(exc).__name__}: {exc}")
+        if COMFY_LOG.exists():
+            lines.append("comfy_log_tail:")
+            lines.extend(COMFY_LOG.read_text(encoding="utf-8", errors="replace").splitlines()[-120:])
+    return "\n".join(lines)
+
 
-    if not torch.cuda.is_available():
-        return "\n".join(lines)
-
-    device = torch.device("cuda:0")
-    lines.append(f"device_name={torch.cuda.get_device_name(0)}")
-    props = torch.cuda.get_device_properties(0)
-    lines.append(f"total_memory_gb={props.total_memory / 1024**3:.2f}")
-
-    tensor = torch.arange(16, device=device, dtype=torch.float32)
-    result = (tensor * 2).sum().item()
-    torch.cuda.synchronize()
-    allocated = torch.cuda.memory_allocated(0) / 1024**2
-    reserved = torch.cuda.memory_reserved(0) / 1024**2
-    lines.append(f"tensor_result={result}")
-    lines.append(f"memory_allocated_mb={allocated:.2f}")
-    lines.append(f"memory_reserved_mb={reserved:.2f}")
+@spaces.GPU(duration=1200)
+def voxcpm_tts_gpu_test() -> str:
+    lines = gpu_status_lines()
+    started = time.time()
+    try:
+        ensure_comfy(lines)
+        prefix = f"voxcpm_tts_gpu_{uuid.uuid4().hex[:8]}"
+        prompt_id = submit_prompt(voxcpm_tts_workflow(prefix))
+        lines.append(f"prompt_id={prompt_id}")
+        history = wait_for_history(prompt_id, timeout=1200)
+        lines.extend(history_summary(history))
+        lines.append(f"elapsed_sec={time.time() - started:.1f}")
+    except Exception as exc:
+        lines.append(f"error={type(exc).__name__}: {exc}")
+        if COMFY_LOG.exists():
+            lines.append("comfy_log_tail:")
+            lines.extend(COMFY_LOG.read_text(encoding="utf-8", errors="replace").splitlines()[-160:])
     return "\n".join(lines)
 
 
 with gr.Blocks(title="VoiceGate") as demo:
     gr.Markdown("# VoiceGate")
-    gr.Markdown("Hugging Face Space deployment scaffold for the VoiceGate ComfyUI workflow.")
-    run = gr.Button("Run GPU smoke test")
-    output = gr.Textbox(label="Status")
-    run.click(fn=gpu_smoke_test, outputs=output)
+    with gr.Row():
+        gpu_run = gr.Button("GPU")
+        comfy_run = gr.Button("ComfyUI")
+        melband_run = gr.Button("MelBand")
+        voxcpm_run = gr.Button("VoxCPM TTS")
+    output = gr.Textbox(label="Status", lines=28)
+    gpu_run.click(fn=gpu_smoke_test, outputs=output)
+    comfy_run.click(fn=comfy_runtime_test, outputs=output)
+    melband_run.click(fn=melband_gpu_test, outputs=output)
+    voxcpm_run.click(fn=voxcpm_tts_gpu_test, outputs=output)
 
 
 if __name__ == "__main__":