app.py

import gc
import hashlib
import json
import os
from typing import Dict, List

import gradio as gr
import spaces

from validator import DeterministicValidator


def _clear_vram():
    import torch

    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()


class DeterministicAudioRenderer:
    def __init__(self, voice: str = "af_bella", lang_code: str = "a", sample_rate: int = 24000):
        self.voice = voice
        self.lang_code = lang_code
        self.sample_rate = sample_rate
        self.pipeline = None

    def _load(self):
        if self.pipeline is None:
            from kokoro import KPipeline

            self.pipeline = KPipeline(lang_code=self.lang_code)

    def render(self, audio_script: str, seed: int, output_path: str = "outputs/audio.wav") -> Dict:
        os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else ".", exist_ok=True)
        text = audio_script.strip()
        if not text:
            return {"status": "error", "reason": "Empty audio_script"}
        self._load()
        import numpy as np
        from scipy.io.wavfile import write as wav_write

        generator = self.pipeline(text, voice=self.voice, speed=1.0)
        segments = []
        seg_count = 0
        for _graphemes, _phonemes, audio_np in generator:
            segments.append(audio_np)
            seg_count += 1
        if not segments:
            return {"status": "error", "reason": "No audio segments generated", "seed": seed}
        full = np.concatenate(segments)
        if full.dtype != np.int16:
            full = (full * 32767).astype(np.int16)
        wav_write(output_path, rate=self.sample_rate, data=full)
        with open(output_path, "rb") as f:
            h = hashlib.sha256(f.read()).hexdigest()[:16]
        return {
            "status": "success",
            "path": output_path,
            "duration_sec": len(full) / self.sample_rate,
            "engine": "kokoro",
            "voice": self.voice,
            "segments": seg_count,
            "hash": h,
            "seed": seed,
        }


class DeterministicVideoRenderer:
    def __init__(self, model_id: str = "Wan-AI/Wan2.1-T2V-14B-Diffusers"):
        self.wan_id = model_id
        self.wan_pipe = None

    def _load(self):
        import torch

        if self.wan_pipe:
            return
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA GPU is required for Wan2.1 video rendering")
        _clear_vram()
        from diffusers import AutoencoderKLWan, WanPipeline
        from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler

        vae = AutoencoderKLWan.from_pretrained(self.wan_id, subfolder="vae", torch_dtype=torch.float32)
        pipe = WanPipeline.from_pretrained(self.wan_id, vae=vae, torch_dtype=torch.bfloat16)
        pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config, flow_shift=5.0)
        pipe.to("cuda")
        self.wan_pipe = pipe

    def _unload(self):
        if self.wan_pipe:
            del self.wan_pipe
            self.wan_pipe = None
            _clear_vram()

    def render(self, media_spec: Dict) -> Dict:
        import torch

        prompt = media_spec["governed_prompt"]
        seed = media_spec["seed"]
        w, h = media_spec["resolution"]
        frames = media_spec["frames"]
        steps = media_spec["steps"]
        guidance = media_spec["guidance_scale"]
        assert (frames - 1) % 4 == 0, f"frames must be 4k+1, got {frames}"
        assert w % 32 == 0 and h % 32 == 0, "resolution must be divisible by 32"
        self._load()
        try:
            generator = torch.Generator(device="cuda").manual_seed(seed)
            result = self.wan_pipe(
                prompt=prompt,
                width=w,
                height=h,
                num_frames=frames,
                num_inference_steps=steps,
                guidance_scale=guidance,
                generator=generator,
            )
            video = result.frames[0]
            output_path = f"outputs/video_{seed}_{frames}f.mp4"
            os.makedirs("outputs", exist_ok=True)
            from diffusers.utils import export_to_video

            export_to_video(video, output_path, fps=media_spec["fps"])
            with open(output_path, "rb") as f:
                hsh = hashlib.sha256(f.read()).hexdigest()[:16]
            return {
                "status": "success",
                "path": output_path,
                "frames": len(video),
                "fps": media_spec["fps"],
                "duration_sec": len(video) / media_spec["fps"],
                "resolution": f"{w}x{h}",
                "seed": seed,
                "engine": self.wan_id,
                "hash": hsh,
            }
        finally:
            self._unload()


class DeterministicMediaRenderer:
    def __init__(self):
        self.audio = DeterministicAudioRenderer(voice="af_bella")
        self.video = DeterministicVideoRenderer()

    def render(self, yt: Dict) -> Dict:
        spec = yt.get("media_spec", {})
        prov = spec.get("provenance", {})
        if not spec or "seed" not in spec:
            return {"status": "REJECTED", "reason": "Missing deterministic media_spec"}
        if not prov.get("source_audit_hash") or not prov.get("reviewer_id"):
            return {"status": "REJECTED", "reason": "Missing provenance. Space2 accepts only human-governed outputs."}
        proposal = yt.get("proposal", {})
        payload = proposal.get("payload", {}) if isinstance(proposal, dict) else {}
        audio_script = spec.get("audio_script", "").strip()
        if not audio_script:
            concepts = payload.get("concepts", [])
            rels = payload.get("inferences", []) or payload.get("proposed_relationships", [])
            audio_script = self._assemble_audio_script(yt, concepts, rels)
        audio_result = self.audio.render(audio_script, spec["seed"], f"outputs/audio_{spec['seed']}.wav")
        if audio_result.get("status") != "success":
            return {"status": "VALIDATION_FAILED", "reason": audio_result.get("reason", "Audio render failed")}
        try:
            video_result = self.video.render(spec)
        except Exception as e:
            return {"status": "VALIDATION_FAILED", "reason": f"Video render failed: {e}"}
        validator = DeterministicValidator()
        v_ok, v_msg = validator.validate_media(spec, video_result, artifact_type="video")
        if not v_ok:
            return {"status": "VALIDATION_FAILED", "reason": v_msg}
        a_ok, a_msg = validator.validate_media(spec, audio_result, artifact_type="audio")
        if not a_ok:
            return {"status": "VALIDATION_FAILED", "reason": a_msg}
        return {
            "status": "DETERMINISTIC_RENDER_COMPLETE",
            "source_audit_hash": prov["source_audit_hash"],
            "reviewer": prov["reviewer_id"],
            "seed": spec["seed"],
            "audio": audio_result,
            "video": video_result,
            "video_validation": v_msg,
            "audio_validation": a_msg,
            "invariant": "NO_WRITE_NO_ADAPTIVE_NO_INTERPRETATION",
        }

    def _assemble_audio_script(self, yt: Dict, concepts: List[str], rels: List[Dict]) -> str:
        scores = (yt.get("alignment", "N/A"), yt.get("adjudication", "N/A"), yt.get("readiness", "N/A"))
        script = f"Governed Knowledge Narration. Alignment {scores[0]}. Adjudication {scores[1]}. Readiness {scores[2]}. "
        script += "Entities: " + ", ".join(concepts) + ". "
        if rels:
            script += "Verified causal relations: "
            for r in rels[:3]:
                if isinstance(r, dict):
                    script += f"{r.get('source')} to {r.get('target')}. "
        script += "Deterministic and governed. All findings passed Markov-Blanket validation and human review."
        return script


renderer = DeterministicMediaRenderer()


@spaces.GPU
def receive_and_render(yt_json: str):
    try:
        payload = json.loads(yt_json)
    except Exception as e:
        return f"REJECTED: Invalid JSON. {e}", None, None, "", ""
    if "yt" in payload:
        yt = payload["yt"]
    elif "alignment" in payload and "media_spec" in payload:
        yt = payload
    else:
        return "REJECTED: Missing yt or media_spec.", None, None, "", ""
    result = renderer.render(yt)
    if result.get("status") == "REJECTED":
        return result["reason"], None, None, "", ""
    if result.get("status") == "VALIDATION_FAILED":
        return f"VALIDATION_FAILED: {result['reason']}", None, None, "", ""
    audio_path = result["audio"]["path"] if result["audio"]["status"] == "success" else None
    video_path = result["video"]["path"] if result["video"]["status"] == "success" else None
    receipt = {
        "status": result["status"],
        "source_audit_hash": result["source_audit_hash"],
        "reviewer": result["reviewer"],
        "seed": result["seed"],
        "audio_hash": result["audio"].get("hash"),
        "video_hash": result["video"].get("hash"),
        "audio_duration": result["audio"].get("duration_sec"),
        "video_duration": result["video"].get("duration_sec"),
        "video_validation": result.get("video_validation"),
        "audio_validation": result.get("audio_validation"),
        "invariant": result["invariant"],
    }
    return json.dumps(receipt, indent=2), audio_path, video_path, json.dumps(result["audio"], indent=2), json.dumps(result["video"], indent=2)


with gr.Blocks(
    title="Deterministic Media Renderer",
    css="""
body{background-color:#050505;color:#e0e0e0;}
.gradio-container{font-family:'Courier New',monospace;}
.tabitem{background-color:#0a0a0a!important;}
button{background-color:#111!important;color:#00ff88!important;border:1px solid #00ff88!important;}
input,textarea{background-color:#111!important;color:#fff!important;border-color:#333!important;}
""",
) as demo:
    gr.Markdown("# DETERMINISTIC MEDIA RENDERER (SPACE 2)\n**RECEIVER ONLY.** Accepts governed `yt` from Bounded Agent (Space1). No adaptive behavior. No writeback. No LLM interpretation.")
    gr.Markdown("### Voice: **af_bella** (locked) | Engine: **Kokoro-TTS** (CPU) | Video: **Wan2.1-T2V-14B** (GPU, unload after)")
    yt_input = gr.TextArea(label="yt payload from Space1 (JSON)", lines=20, placeholder="Paste complete yt JSON from Space1 EXPORT tab")
    render_btn = gr.Button("DETERMINISTIC RENDER (Audio + Video)")
    receipt = gr.Code(label="Render Receipt + Provenance", language="json")
    with gr.Row():
        audio_out = gr.Audio(label="Governed Audio (af_bella)", autoplay=False)
        video_out = gr.Video(label="Governed Video (Wan2.1 14B)")
    with gr.Row():
        audio_meta = gr.Code(label="Audio Metadata", language="json")
        video_meta = gr.Code(label="Video Metadata", language="json")
    render_btn.click(fn=receive_and_render, inputs=yt_input, outputs=[receipt, audio_out, video_out, audio_meta, video_meta], api_name=False)
    gr.HTML('<div style="color:#ff5555;font-family:monospace;margin-top:20px;padding:12px;border:1px solid #330000;"><b>INVARIANTS ENFORCED:</b><br>Audio voice locked: af_bella. No variation. No temperature.<br>Video seed locked from source audit hash - reproducible<br>Prompt verbatim from governed script - no LLM enhancement<br>Parameters policy-fixed - no deviation<br>No writeback to Space1 - read-only receiver<br>Audio CPU (Kokoro). Video GPU (Wan2.1). Isolated resources.<br>Output hashed - provenance chain intact<br>Validator checks deterministic integrity before release</div>')


if __name__ == "__main__":
    demo.launch()