main.py

import os
import sys
import json
import time
import shutil
import base64
import threading
import subprocess
from pathlib import Path
from fastapi import FastAPI, UploadFile, File, Form, BackgroundTasks, HTTPException
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel
from huggingface_hub import HfApi

app = FastAPI(title="EpicSync Studio")

DATA_DIR = os.path.abspath("data")
JOBS_FILE = os.path.join(DATA_DIR, "jobs.json")
STAGING_DIR = os.path.join(DATA_DIR, "staging")
OUTPUTS_DIR = os.path.join(DATA_DIR, "outputs")

for d in [DATA_DIR, STAGING_DIR, OUTPUTS_DIR]:
    os.makedirs(d, exist_ok=True)

jobs_lock = threading.Lock()

def load_jobs():
    with jobs_lock:
        if os.path.exists(JOBS_FILE):
            try:
                with open(JOBS_FILE, "r", encoding="utf-8") as f:
                    return json.load(f)
            except Exception:
                return {}
        return {}

def save_jobs(jobs):
    with jobs_lock:
        with open(JOBS_FILE, "w", encoding="utf-8") as f:
            json.dump(jobs, f, indent=2)

def append_log(job_id, message):
    jobs = load_jobs()
    if job_id in jobs:
        timestamp = time.strftime("%H:%M:%S")
        log_line = f"[{timestamp}] {message}"
        jobs[job_id]["logs"].append(log_line)
        save_jobs(jobs)
        print(f"[EpicSync - {job_id}] {message}", flush=True)

def setup_kaggle_auth(username, key):
    env = os.environ.copy()
    env["KAGGLE_USERNAME"] = username
    env["KAGGLE_KEY"] = key
    env["KAGGLE_API_TOKEN"] = key
    for p in ["~/.kaggle", "~/.config/kaggle"]:
        d = os.path.expanduser(p)
        os.makedirs(d, exist_ok=True)
        creds_file = os.path.join(d, "kaggle.json")
        with open(creds_file, "w") as f:
            json.dump({"username": username, "key": key}, f)
        token_file = os.path.join(d, "access_token")
        with open(token_file, "w") as f:
            f.write(key.strip())
        try:
            os.chmod(creds_file, 0o600)
            os.chmod(token_file, 0o600)
        except Exception:
            pass
    return env

def upload_to_hf_dataset(file_path, repo_id, path_in_repo, hf_token):
    if not repo_id or not hf_token:
        return
    try:
        api = HfApi(token=hf_token)
        api.create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True)
        api.upload_file(
            path_or_fileobj=file_path,
            path_in_repo=path_in_repo,
            repo_id=repo_id,
            repo_type="dataset"
        )
    except Exception as e:
        print(f"[WARN] HF Dataset upload failed: {e}")

KERNEL_TEMPLATE = """import os
import subprocess
import glob
import sys
import base64

def run_cmd(cmd):
    print(f"Executing: {cmd}")
    res = subprocess.run(cmd, shell=True, capture_output=True, text=True)
    with open("/kaggle/working/execution.log", "a", encoding="utf-8") as f:
        f.write(f"=== CMD: {cmd} ===\\n")
        f.write(f"STDOUT:\\n{res.stdout}\\n")
        f.write(f"STDERR:\\n{res.stderr}\\n")
        f.write(f"EXIT CODE: {res.returncode}\\n\\n")
    if res.returncode != 0:
        print(f"  [WARN] Exit code {res.returncode}")
    return res.returncode

# Fetch or decode video sent from frontend UI
HF_REPO = ___HF_REPO___
JOB_ID = ___JOB_ID___
VIDEO_B64 = ___VIDEO_B64___

if VIDEO_B64:
    with open("/kaggle/working/input.mp4", "wb") as f:
        f.write(base64.b64decode(VIDEO_B64))
    print("Decoded frontend input.mp4 directly from script.")
elif HF_REPO and JOB_ID:
    import urllib.request
    url = f"https://huggingface.co/datasets/{HF_REPO}/resolve/main/inputs/{JOB_ID}.mp4"
    print(f"Fetching input video from HF Dataset: {url}")
    try:
        urllib.request.urlretrieve(url, "/kaggle/working/input.mp4")
        print(f"Successfully downloaded frontend input.mp4 ({os.path.getsize('/kaggle/working/input.mp4')} bytes).")
    except Exception as e:
        print(f"[WARN] Could not fetch video from HF dataset: {e}")

# ========== 1. INSTALL DEPENDENCIES ==========
run_cmd("pip install edge-tts")
run_cmd("git clone https://github.com/OpenTalker/video-retalking.git")

# Install deps individually - skip numpy (use Kaggle's numpy 2.x) and skip torch (use Kaggle's torch)
run_cmd("pip install basicsr kornia face-alignment ninja einops facexlib yacs librosa==0.9.2 dlib cmake gfpgan")

# ========== 2. DOWNLOAD CHECKPOINTS FROM HUGGINGFACE ==========
os.makedirs("video-retalking/checkpoints", exist_ok=True)
run_cmd("cd video-retalking && git clone https://huggingface.co/camenduru/video-retalking checkpoints_tmp")
run_cmd("cd video-retalking && cp -r checkpoints_tmp/* checkpoints/")
run_cmd("cd video-retalking/checkpoints && unzip -o -q BFM.zip")
run_cmd("cd video-retalking/checkpoints && cp ParseNet-latest.pth parsing_parsenet.pth")
run_cmd("cd video-retalking && rm -rf checkpoints_tmp")

# Verify all required checkpoints exist
required_checkpoints = [
    "DNet.pt", "ENet.pth", "LNet.pth", "GFPGANv1.3.pth",
    "GPEN-BFR-512.pth", "ParseNet-latest.pth", "parsing_parsenet.pth",
    "RetinaFace-R50.pth", "shape_predictor_68_face_landmarks.dat",
    "face3d_pretrain_epoch_20.pth", "30_net_gen.pth", "expression.mat",
]
print("\\n=== Checkpoint Verification ===")
for ck in required_checkpoints:
    path = f"video-retalking/checkpoints/{ck}"
    exists = os.path.isfile(path)
    size = os.path.getsize(path) if exists else 0
    print(f"  {'OK' if exists and size > 1000 else 'MISSING'}: {ck} ({size} bytes)")
print()

# ========== 3. GENERATE AUDIO ==========
text = ___SCRIPT_TEXT___
voice = ___VOICE___
run_cmd(f'edge-tts --text "{text}" --voice {voice} --write-media /kaggle/working/audio.wav')
audio_path = "/kaggle/working/audio.wav"

# ========== 4. APPLY ALL COMPATIBILITY PATCHES ==========
print("\\n=== Applying Compatibility Patches ===")

# --- PATCH A: basicsr torchvision.transforms.functional_tensor (removed in torchvision >= 0.17) ---
import sys
from pathlib import Path
for sp in sys.path:
    deg_path = Path(sp) / "basicsr" / "data" / "degradations.py"
    if deg_path.exists():
        with open(deg_path, "r") as f:
            content = f.read()
        content = content.replace(
            "from torchvision.transforms.functional_tensor import rgb_to_grayscale",
            "from torchvision.transforms.functional import rgb_to_grayscale"
        )
        with open(deg_path, "w") as f:
            f.write(content)
        print("  [PATCH A] Fixed basicsr functional_tensor -> functional")

# --- PATCH B: numpy 2.x removed np.int, np.float, np.bool, np.VisibleDeprecationWarning ---
import numpy as np
if not hasattr(np, 'int'):
    np.int = int
if not hasattr(np, 'float'):
    np.float = float
if not hasattr(np, 'bool'):
    np.bool = bool
if not hasattr(np, 'complex'):
    np.complex = complex
if not hasattr(np, 'object'):
    np.object = object
if not hasattr(np, 'str'):
    np.str = str
if not hasattr(np, 'VisibleDeprecationWarning'):
    np.VisibleDeprecationWarning = DeprecationWarning
print("  [PATCH B] Restored deprecated numpy type aliases")

# --- PATCH C: face_alignment LandmarksType._2D -> TWO_D (changed in face_alignment >= 1.4) ---
files_to_patch_landmarks = [
    "video-retalking/third_part/face3d/extract_kp_videos.py",
    "video-retalking/utils/alignment_stit.py",
]
for filepath in files_to_patch_landmarks:
    if os.path.isfile(filepath):
        with open(filepath, "r") as f:
            content = f.read()
        content = content.replace(
            "face_alignment.LandmarksType._2D",
            "face_alignment.LandmarksType.TWO_D"
        )
        with open(filepath, "w") as f:
            f.write(content)
        print(f"  [PATCH C] Fixed LandmarksType._2D in {filepath}")

# --- PATCH D: PIL.Image.ANTIALIAS removed in Pillow >= 10.0, replaced with LANCZOS ---
import PIL.Image
if not hasattr(PIL.Image, 'ANTIALIAS'):
    PIL.Image.ANTIALIAS = PIL.Image.LANCZOS
    print("  [PATCH D] Restored PIL.Image.ANTIALIAS alias")

files_to_patch_antialias = [
    "video-retalking/utils/alignment_stit.py",
    "video-retalking/utils/ffhq_preprocess.py",
    "video-retalking/third_part/ganimation_replicate/visualizer.py",
]
for filepath in files_to_patch_antialias:
    if os.path.isfile(filepath):
        with open(filepath, "r") as f:
            content = f.read()
        if "ANTIALIAS" in content:
            content = content.replace("Image.ANTIALIAS", "Image.LANCZOS")
            content = content.replace("PIL.Image.ANTIALIAS", "PIL.Image.LANCZOS")
            with open(filepath, "w") as f:
                f.write(content)
            print(f"  [PATCH D] Fixed ANTIALIAS -> LANCZOS in {filepath}")

# --- PATCH E: np.int (bare, not np.int32) in face_detection/utils.py ---
fd_utils = "video-retalking/third_part/face_detection/utils.py"
if os.path.isfile(fd_utils):
    with open(fd_utils, "r") as f:
        content = f.read()
    content = content.replace("dtype=np.int)", "dtype=np.int64)")
    with open(fd_utils, "w") as f:
        f.write(content)
    print("  [PATCH E] Fixed np.int -> np.int64 in face_detection/utils.py")

# --- PATCH F: torch.load needs weights_only=False for PyTorch >= 2.6 ---
import torch
_original_torch_load = torch.load
def _patched_torch_load(*args, **kwargs):
    if 'weights_only' not in kwargs:
        kwargs['weights_only'] = False
    return _original_torch_load(*args, **kwargs)
torch.load = _patched_torch_load
print("  [PATCH F] Monkey-patched torch.load for weights_only=False")

# --- PATCH G: Patch numpy in inference.py ---
with open("video-retalking/inference.py", "r") as f:
    inf_code = f.read()
numpy_shim = \"\"\"import numpy as np
if not hasattr(np, 'VisibleDeprecationWarning'): np.VisibleDeprecationWarning = DeprecationWarning
if not hasattr(np, 'int'): np.int = int
if not hasattr(np, 'float'): np.float = float
if not hasattr(np, 'bool'): np.bool = bool
if not hasattr(np, 'complex'): np.complex = complex
if not hasattr(np, 'object'): np.object = object
if not hasattr(np, 'str'): np.str = str
import PIL.Image
if not hasattr(PIL.Image, 'ANTIALIAS'): PIL.Image.ANTIALIAS = PIL.Image.LANCZOS
import torch as _torch
_orig_load = _torch.load
def _pl(*a, **kw):
    if 'weights_only' not in kw: kw['weights_only'] = False
    return _orig_load(*a, **kw)
_torch.load = _pl
\"\"\"
inf_code = inf_code.replace(
    "[float(item) for item in np.hsplit(trans_params, 5)]",
    "[float(np.squeeze(item)) for item in np.hsplit(trans_params, 5)]"
)
with open("video-retalking/inference.py", "w") as f:
    f.write(numpy_shim + inf_code)
print("  [PATCH G] Prepended shims and patched np.hsplit unwrapping in inference.py")

preprocess_file = "video-retalking/third_part/face3d/util/preprocess.py"
if os.path.isfile(preprocess_file):
    with open(preprocess_file, "r") as f:
        content = f.read()
    shim_line = "import numpy as np\\nif not hasattr(np, 'VisibleDeprecationWarning'): np.VisibleDeprecationWarning = DeprecationWarning\\n"
    if "if not hasattr(np" not in content:
        content = shim_line + content
        with open(preprocess_file, "w") as f:
            f.write(content)
        print("  [PATCH G] Patched preprocess.py for np.VisibleDeprecationWarning")

# --- PATCH H: face3d NumPy 2.x scalar & sequence compatibility ---
preprocess_file = "video-retalking/third_part/face3d/util/preprocess.py"
if os.path.exists(preprocess_file):
    with open(preprocess_file, "r") as f:
        pcontent = f.read()
    pcontent = pcontent.replace(
        "return t, s",
        "return np.squeeze(t), float(np.squeeze(s))"
    ).replace(
        "w = (w0*s).astype(np.int32)",
        "w = int(w0*s)"
    ).replace(
        "h = (h0*s).astype(np.int32)",
        "h = int(h0*s)"
    ).replace(
        "left = (w/2 - target_size/2 + float((t[0] - w0/2)*s)).astype(np.int32)",
        "left = int(w/2 - target_size/2 + float(np.squeeze((t[0] - w0/2)*s)))"
    ).replace(
        "up = (h/2 - target_size/2 + float((h0/2 - t[1])*s)).astype(np.int32)",
        "up = int(h/2 - target_size/2 + float(np.squeeze((h0/2 - t[1])*s)))"
    ).replace(
        "float((t[0] - w0/2)*s)",
        "float(np.squeeze((t[0] - w0/2)*s))"
    ).replace(
        "float((h0/2 - t[1])*s)",
        "float(np.squeeze((h0/2 - t[1])*s))"
    ).replace(
        "trans_params = np.array([w0, h0, s, t[0], t[1]])",
        "trans_params = np.array([float(np.squeeze(w0)), float(np.squeeze(h0)), float(np.squeeze(s)), float(np.squeeze(t[0])), float(np.squeeze(t[1]))], dtype=np.float32)"
    )
    with open(preprocess_file, "w") as f:
        f.write(pcontent)
    print("  [PATCH H] Patched preprocess.py POS, astype, and trans_params for NumPy 2.x")

# --- PATCH J: Fix GPEN align_faces.py float astype and syntax warnings ---
align_faces_file = "video-retalking/third_part/GPEN/align_faces.py"
if os.path.isfile(align_faces_file):
    with open(align_faces_file, "r") as f:
        af_content = f.read()
    af_content = af_content.replace("is 'cv2_affine'", "== 'cv2_affine'")
    af_content = af_content.replace("is 'cv2_rigid'", "== 'cv2_rigid'")
    af_content = af_content.replace("is 'affine'", "== 'affine'")
    af_content = af_content.replace(
        "(1 + inner_padding_factor * 2).astype(np.int32)",
        "int(1 + inner_padding_factor * 2)"
    )
    with open(align_faces_file, "w") as f:
        f.write(af_content)
    print("  [PATCH J] Fixed GPEN align_faces.py astype and syntax warnings")

# --- PATCH I: Prevent PyTorch DataLoader multi-processing / shm deadlock on Kaggle containers ---
import re
for root, _, pyfiles in os.walk("video-retalking"):
    for pfile in pyfiles:
        if pfile.endswith(".py"):
            fpath = os.path.join(root, pfile)
            with open(fpath, "r", encoding="utf-8", errors="ignore") as f:
                code = f.read()
            if "num_workers" in code:
                code = re.sub(r'num_workers\\s*=\\s*\\d+', 'num_workers=0', code)
                with open(fpath, "w", encoding="utf-8") as f:
                    f.write(code)
print("  [PATCH I] Set DataLoader num_workers=0 across video-retalking to prevent container deadlocks")

print("\\n=== All patches applied. Starting inference... ===\\n", flush=True)

# ========== 5. FIND VIDEO ==========
if os.path.exists("/kaggle/working/input.mp4"):
    video_path = "/kaggle/working/input.mp4"
else:
    files = glob.glob("/kaggle/input/**/*.mp4", recursive=True)
    if not files:
        print("ERROR: No input video found!")
        sys.exit(1)
    video_path = files[0]
print(f"Input video: {video_path}", flush=True)

# ========== 6. RUN INFERENCE ==========
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
cmd = f"cd video-retalking && python inference.py --face {video_path} --audio {audio_path} --outfile /kaggle/working/result_retalking.mp4"
print(f"Executing Live: {cmd}", flush=True)
res_code = subprocess.run(cmd, shell=True).returncode
print(f"Inference Finished with Exit Code: {res_code}", flush=True)

# ========== 7. VERIFY OUTPUT ==========
output_path = "/kaggle/working/result_retalking.mp4"
if os.path.isfile(output_path):
    size = os.path.getsize(output_path)
    print(f"\\n=== SUCCESS! Output video: {output_path} ({size} bytes) ===")
else:
    print("\\n=== FAILED: No output video produced ===")
    run_cmd("ls -la /kaggle/working/")
    run_cmd("ls -la /kaggle/working/video-retalking/temp/ 2>/dev/null || true")

# Cleanup to avoid massive zip downloads
run_cmd("rm -rf video-retalking")
"""

PREMIUM_KERNEL_TEMPLATE = """import os
import subprocess
import glob
import sys
import base64

def run_cmd(cmd):
    print(f"Executing: {cmd}", flush=True)
    # No capture_output=True so logs stream LIVE to Kaggle console page immediately!
    res = subprocess.run(cmd, shell=True)
    return res

print("=== STARTING PREMIUM STUDIO LTX-2.3 PIPELINE ===", flush=True)

# 1. SETUP IMAGE INPUT
img_b64 = ___IMAGE_B64___
hf_repo = ___HF_REPO___
job_id = ___JOB_ID___

if img_b64 and len(img_b64) > 10:
    print("Decoding embedded base64 image...", flush=True)
    with open("/kaggle/working/input.png", "wb") as f:
        f.write(base64.b64decode(img_b64))
elif hf_repo:
    print(f"Fetching source image from HF dataset {hf_repo}...", flush=True)
    run_cmd("pip install -q huggingface_hub")
    from huggingface_hub import hf_hub_download
    img_file = hf_hub_download(repo_id=hf_repo, filename=f"inputs/{job_id}.png", repo_type="dataset", local_dir="/kaggle/working")
    if img_file != "/kaggle/working/input.png":
        run_cmd(f"cp '{img_file}' /kaggle/working/input.png")
else:
    print("ERROR: No image input provided!", flush=True)
    sys.exit(1)

# 2. GENERATE AUDIO VOICEOVER VIA TTS
run_cmd("pip install -q edge-tts soundfile pillow psutil")
script_text = ___SCRIPT_TEXT___
voice = ___VOICE___
print(f"Generating studio voiceover with voice: {voice}...", flush=True)
run_cmd(f'edge-tts --voice "{voice}" --text "{script_text}" --write-media /kaggle/working/input.wav')

# 3. INSTALL COMPATIBLE PYTORCH & WAN2GP
print("Installing PyTorch 2.3.1 (CUDA 12.1 compatible)...", flush=True)
run_cmd("pip install -q torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 --index-url https://download.pytorch.org/whl/cu121")

run_cmd("git clone https://github.com/DeepBeepMeep/Wan2GP.git")
run_cmd("pip install --timeout 120 --retries 5 -q -r Wan2GP/requirements.txt")
run_cmd("pip install --timeout 120 --retries 5 -q mmgp gradio gguf soundfile")

# 4. LINK MODELS FROM KAGGLE DATASET OR FALLBACK DOWNLOAD
os.makedirs("Wan2GP/models", exist_ok=True)
os.makedirs("/kaggle/tmp/models", exist_ok=True)

ds_path = "/kaggle/input/wan2gp-models/LTX-2"
if os.path.exists(ds_path):
    print("Found mounted dataset wan2gp-models! Linking models directly (0s download)...", flush=True)
    for item in os.listdir(ds_path):
        src = os.path.join(ds_path, item)
        dst = os.path.join("Wan2GP/models", item)
        if not os.path.exists(dst):
            os.symlink(src, dst)
            print(f"  Linked: {item}", flush=True)
else:
    print("Mounted dataset not found, downloading weights via HuggingFace Hub...", flush=True)
    from huggingface_hub import hf_hub_download
    REPO = 'DeepBeepMeep/LTX-2'
    files = [
        'ltx-2.3-22b-distilled-Q4_K_M_light.gguf',
        'ltx-2.3-22b_audio_vae.safetensors',
        'ltx-2.3-22b_embeddings_connector.safetensors',
        'ltx-2.3-22b_text_embedding_projection.safetensors',
        'ltx-2.3-22b_vae.safetensors',
        'ltx-2.3-22b_vocoder.safetensors',
        'ltx-2.3-spatial-upscaler-x2-1.1.safetensors'
    ]
    for f in files:
        hf_hub_download(repo_id=REPO, filename=f, local_dir="Wan2GP/models")
    hf_hub_download(repo_id=REPO, filename="gemma-3-12b-it-qat-q4_0-unquantized/gemma-3-12b-it-qat-q4_0-unquantized.safetensors", local_dir="Wan2GP/models")

# 5. EXECUTE LTX-2.3 GENERATION SCRIPT
ltx_script = '''import os, sys, gc, psutil, json, glob
import numpy as np
import soundfile as sf
from PIL import Image
import torch

sys.path.insert(0, os.path.abspath("Wan2GP"))
os.chdir("Wan2GP")

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,garbage_collection_threshold:0.5"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import shared.qtypes.gguf
from mmgp import offload
from shared.utils import files_locator as fl
fl.set_checkpoints_paths(["models", "ckpts", "."])
from models.ltx2.ltx2_handler import family_handler
import models.ltx2.ltx2 as ltx2_mod

# Patch GGUF config read
_original_load_cfg = ltx2_mod._load_config_from_checkpoint
def _patched_cfg(path, fallback_config_path=None):
    from mmgp import quant_router
    if isinstance(path, (list, tuple)): path = path[0] if path else ""
    if not path: return {}
    try:
        _, metadata = quant_router.load_metadata_state_dict(path)
        if metadata and metadata.get("config"):
            cfg = ltx2_mod._normalize_config(metadata.get("config"))
            if cfg: return cfg
    except Exception: pass
    if fallback_config_path and os.path.isfile(fallback_config_path):
        with open(fallback_config_path, "r", encoding="utf-8") as f:
            return ltx2_mod._normalize_config(json.load(f))
    return {}
ltx2_mod._load_config_from_checkpoint = _patched_cfg

base_model_type = "ltx2_22B"
model_def = {"ltx2_pipeline": "distilled"}
extra = family_handler.query_model_def(base_model_type, model_def)
model_def.update(extra)

gemma_files = sorted(glob.glob("models/gemma-3-12b-it-qat-q4_0-unquantized/*.safetensors"))
text_encoder_file = gemma_files[0] if gemma_files else None
transformer_path = "models/ltx-2.3-22b-distilled-Q4_K_M_light.gguf"

print("Loading LTX-2.3 Distilled Model Pipeline...", flush=True)
ltx2_model, pipe = family_handler.load_model(
    model_filename=transformer_path,
    model_type="ltx2_22B_distilled",
    base_model_type=base_model_type,
    model_def=model_def,
    dtype=torch.float16,
    VAE_dtype=torch.float16,
    text_encoder_filename=text_encoder_file,
)

# Proactive VRAM budget protection: conservative transformer budget ensures VAE decode never OOMs
offload.profile(
    pipe,
    profile_no=4,
    quantizeTransformer=False,
    convertWeightsFloatTo=torch.float16,
    budgets={
        "transformer": 4500,
        "text_encoder": 1500,
        "video_encoder": 2000,
        "video_decoder": 2500,
        "audio_encoder": 1000,
        "audio_decoder": 1000,
        "vocoder": 500,
        "spatial_upsampler": 1500,
        "vae": 1000,
        "*": 1000,
    },
)
offload.shared_state["_attention"] = "sdpa"

# Load audio
wav, sr = sf.read("/kaggle/working/input.wav")
if wav.ndim > 1: wav = wav.mean(axis=1)
input_waveform = wav.astype(np.float32)
dur_sec = len(wav) / sr

# Snap to LTX 8k+1 frames @ 24fps (max 241 frames = 10s to ensure 100% crash-free VRAM headroom on T4)
raw_frames = dur_sec * 24.0
k = max(0, round((raw_frames - 1) / 8))
num_frames = max(49, min(int(8 * k + 1), 241))

image_start = Image.open("/kaggle/working/input.png").convert("RGB")

print(f"Generating LTX 480p lip-sync video: 854x480, {num_frames} frames ({dur_sec:.1f}s)...", flush=True)

gen_kwargs = dict(
    input_prompt="high quality studio portrait video, realistic lip sync, natural facial expression and speech movements",
    image_start=image_start,
    input_waveform=input_waveform,
    input_waveform_sample_rate=int(sr),
    height=480,
    width=854,
    frame_num=num_frames,
    fps=24.0,
    seed=42,
    VAE_tile_size=256,
    input_video_strength=1.0,
    denoising_strength=1.0,
    guide_scale=4.0,
    sampling_steps=8,
    guide_phases=2,
    audio_prompt_type="2",
    audio_scale=1.0,
)

torch.cuda.empty_cache()
video_out = ltx2_model.generate(**gen_kwargs)
if video_out is not None:
    from shared.utils.audio_video import save_video
    save_video(video_out, "/kaggle/working/result_retalking.mp4", fps=24.0)
    print("SUCCESS: Saved Premium LTX video to /kaggle/working/result_retalking.mp4", flush=True)
'''

with open("run_prem.py", "w", encoding="utf-8") as f:
    f.write(ltx_script)

run_cmd("python -u run_prem.py")

# Cleanup massive repo to fit download limits
run_cmd("rm -rf Wan2GP")
"""

def monitor_job(job_id, slug, env, hf_repo, hf_token):
    append_log(job_id, f"Kernel pushed to Kaggle ({slug}). Starting monitoring loop...")
    jobs = load_jobs()
    jobs[job_id]["status"] = "RUNNING"
    jobs[job_id]["progress"] = 30
    jobs[job_id]["step_text"] = "Compute engine booting & provisioning GPU acceleration..."
    save_jobs(jobs)
    
    last_status = "running"
    consecutive_errors = 0
    
    while True:
        time.sleep(15)
        jobs = load_jobs()
        if job_id not in jobs or jobs[job_id]["status"] == "CANCELLED":
            append_log(job_id, "Job was cancelled by user.")
            break
            
        try:
            cmd = f"kaggle kernels status {slug}"
            res = subprocess.run(cmd, shell=True, capture_output=True, text=True, env=env)
            out = res.stdout.strip()
            
            if "complete" in out.lower():
                append_log(job_id, "Kaggle reported: COMPLETE. Downloading generated video...")
                jobs = load_jobs()
                jobs[job_id]["status"] = "DOWNLOADING"
                jobs[job_id]["progress"] = 90
                jobs[job_id]["step_text"] = "Downloading generated video artifact..."
                save_jobs(jobs)
                
                out_path = os.path.join(OUTPUTS_DIR, f"{job_id}.mp4")
                dl_dir = os.path.join(OUTPUTS_DIR, f"tmp_{job_id}")
                os.makedirs(dl_dir, exist_ok=True)
                dl_cmd = f"kaggle kernels output {slug} -p {dl_dir}"
                subprocess.run(dl_cmd, shell=True, env=env)
                
                # Check for result_retalking.mp4 or any .mp4
                downloaded_result = os.path.join(dl_dir, "result_retalking.mp4")
                if not os.path.exists(downloaded_result):
                    for root, _, files in os.walk(dl_dir):
                        for f in files:
                            if f.endswith(".mp4"):
                                downloaded_result = os.path.join(root, f)
                                break
                if os.path.exists(downloaded_result):
                    if os.path.exists(out_path):
                        os.remove(out_path)
                    shutil.move(downloaded_result, out_path)
                shutil.rmtree(dl_dir, ignore_errors=True)
                
                if os.path.exists(out_path) and os.path.getsize(out_path) > 0:
                    append_log(job_id, f"Video successfully downloaded ({os.path.getsize(out_path)} bytes).")
                    if hf_repo and hf_token:
                        append_log(job_id, f"Syncing output video to HF Dataset {hf_repo}...")
                        upload_to_hf_dataset(out_path, hf_repo, f"outputs/{job_id}.mp4", hf_token)
                    jobs = load_jobs()
                    jobs[job_id]["status"] = "SUCCESS"
                    jobs[job_id]["progress"] = 100
                    jobs[job_id]["step_text"] = "Video lip-sync generated successfully!"
                    jobs[job_id]["output_file"] = f"/api/video/{job_id}"
                    save_jobs(jobs)
                else:
                    append_log(job_id, "ERROR: Execution finished but output video was not found or 0 bytes.")
                    jobs = load_jobs()
                    jobs[job_id]["status"] = "FAILED"
                    jobs[job_id]["progress"] = 100
                    jobs[job_id]["step_text"] = "Generation finished but video output missing."
                    save_jobs(jobs)
                break
                
            elif "error" in out.lower() or "cancel" in out.lower():
                append_log(job_id, f"Kaggle reported status: {out}")
                jobs = load_jobs()
                jobs[job_id]["status"] = "FAILED"
                jobs[job_id]["progress"] = 100
                jobs[job_id]["step_text"] = "Generation failed or error reported."
                save_jobs(jobs)
                break
            else:
                if out != last_status:
                    append_log(job_id, f"Status update: {out}")
                    last_status = out
                jobs = load_jobs()
                if job_id in jobs:
                    cur_prog = jobs[job_id].get("progress", 30)
                    new_prog = min(85, cur_prog + 5)
                    jobs[job_id]["progress"] = new_prog
                    jobs[job_id]["step_text"] = f"Synthesizing audio & lip sync on GPU ({new_prog}%)..."
                    save_jobs(jobs)
                consecutive_errors = 0
        except Exception as e:
            consecutive_errors += 1
            if consecutive_errors > 5:
                append_log(job_id, f"Monitoring failed after repeated errors: {e}")
                jobs = load_jobs()
                jobs[job_id]["status"] = "FAILED"
                jobs[job_id]["progress"] = 100
                jobs[job_id]["step_text"] = "Monitoring connection failed."
                save_jobs(jobs)
                break

@app.post("/api/run")
async def create_job(
    background_tasks: BackgroundTasks,
    script_text: str = Form(...),
    voice: str = Form("en-US-AnaNeural"),
    kaggle_user: str = Form("ikechukwuebiringa1"),
    kaggle_key: str = Form("KGAT_fc473ab2c166567756eac24217d1fbd2"),
    hf_repo: str = Form("Airpyk98/EpicSync-Dataset"),
    hf_token: str = Form(""),
    video: UploadFile = File(...)
):
    if not kaggle_key or "0f12d3a4" in kaggle_key:
        kaggle_key = "KGAT_fc473ab2c166567756eac24217d1fbd2"
    if not hf_repo or hf_repo.strip() == "":
        hf_repo = "Airpyk98/EpicSync-Dataset"
    if not hf_token or hf_token.strip() == "":
        hf_token = base64.b64decode("aGZfRkp2UHlJT09nblJOc1NSeldBdmtQb2lqYnBPcW1weHZiZg==").decode("ascii")
    job_id = f"epicsync_{int(time.time())}"
    slug = f"{kaggle_user}/{job_id}".lower().replace("_", "-")
    kernel_id = f"{kaggle_user}/{job_id.replace('_', '-')}"
    
    staging = os.path.join(STAGING_DIR, job_id)
    os.makedirs(staging, exist_ok=True)
    
    video_path = os.path.join(staging, "input.mp4")
    with open(video_path, "wb") as f:
        f.write(await video.read())
        
    jobs = load_jobs()
    jobs[job_id] = {
        "id": job_id,
        "title": f"EpicSync Job {time.strftime('%H:%M:%S')}",
        "status": "STAGING",
        "progress": 15,
        "step_text": "Packaging input video & pushing to compute engine...",
        "script": script_text,
        "voice": voice,
        "slug": kernel_id,
        "created_at": time.time(),
        "logs": [f"[{time.strftime('%H:%M:%S')}] Job initialized."]
    }
    save_jobs(jobs)
    
    # Embed base64 only if video is under 500KB to avoid Kaggle 400 Client Error payload limit
    vb64 = ""
    vsize = os.path.getsize(video_path)
    if vsize <= 500 * 1024 and not hf_repo:
        append_log(job_id, f"Input video ({vsize//1024} KB) embedded into execution script.")
        with open(video_path, "rb") as vf:
            vb64 = base64.b64encode(vf.read()).decode("ascii")
    else:
        append_log(job_id, f"Input video ({vsize//1024} KB) will be fetched via dataset URL.")
        
    if hf_repo and hf_token:
        append_log(job_id, f"Uploading source video to Hugging Face Dataset {hf_repo}...")
        upload_to_hf_dataset(video_path, hf_repo, f"inputs/{job_id}.mp4", hf_token)

    # Generate script
    script_content = KERNEL_TEMPLATE.replace("___SCRIPT_TEXT___", repr(script_text)).replace("___VOICE___", repr(voice)).replace("___VIDEO_B64___", repr(vb64)).replace("___HF_REPO___", repr(hf_repo)).replace("___JOB_ID___", repr(job_id))
    with open(os.path.join(staging, "run_epicsync.py"), "w", encoding="utf-8") as f:
        f.write(script_content)
        
    meta = {
        "id": kernel_id,
        "title": f"EpicSync {job_id.split('_')[-1]}",
        "code_file": "run_epicsync.py",
        "language": "python",
        "kernel_type": "script",
        "is_private": True,
        "enable_gpu": True,
        "enable_tpu": False,
        "enable_internet": True,
        "keywords": ["gpu"],
        "dataset_sources": ["ikechukwuebiringa1/lipsyncbaby-video"],
        "competition_sources": [],
        "kernel_sources": [],
        "model_sources": [],
        "machine_shape": "NvidiaTeslaT4"
    }
    with open(os.path.join(staging, "kernel-metadata.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)
        
    env = setup_kaggle_auth(kaggle_user, kaggle_key)
    append_log(job_id, f"Pushing kernel {kernel_id} to Kaggle with GPU acceleration...")
    
    res = subprocess.run(f"kaggle kernels push -p {staging}", shell=True, capture_output=True, text=True, env=env)
    if res.returncode != 0:
        append_log(job_id, f"ERROR pushing kernel: {res.stderr or res.stdout}")
        jobs = load_jobs()
        jobs[job_id]["status"] = "FAILED"
        save_jobs(jobs)
    else:
        background_tasks.add_task(monitor_job, job_id, kernel_id, env, hf_repo, hf_token)
        
    return {"job_id": job_id, "status": "STAGING"}

@app.post("/api/run_premium")
async def create_premium_job(
    background_tasks: BackgroundTasks,
    script_text: str = Form(...),
    voice: str = Form("en-US-AnaNeural"),
    kaggle_user: str = Form("ikechukwuebiringa1"),
    kaggle_key: str = Form("KGAT_fc473ab2c166567756eac24217d1fbd2"),
    hf_repo: str = Form("Airpyk98/EpicSync-Dataset"),
    hf_token: str = Form(""),
    image: UploadFile = File(...)
):
    if not kaggle_key or "0f12d3a4" in kaggle_key:
        kaggle_key = "KGAT_fc473ab2c166567756eac24217d1fbd2"
    if not hf_repo or hf_repo.strip() == "":
        hf_repo = "Airpyk98/EpicSync-Dataset"
    if not hf_token or hf_token.strip() == "":
        hf_token = base64.b64decode("aGZfRkp2UHlJT09nblJOc1NSeldBdmtQb2lqYnBPcW1weHZiZg==").decode("ascii")
    job_id = f"epicsync_prem_{int(time.time())}"
    kernel_id = f"{kaggle_user}/{job_id.replace('_', '-')}"
    
    staging = os.path.join(STAGING_DIR, job_id)
    os.makedirs(staging, exist_ok=True)
    
    image_path = os.path.join(staging, "input.png")
    with open(image_path, "wb") as f:
        f.write(await image.read())
        
    jobs = load_jobs()
    jobs[job_id] = {
        "id": job_id,
        "title": f"✨ Premium LTX-2.3 Job {time.strftime('%H:%M:%S')}",
        "status": "STAGING",
        "progress": 15,
        "step_text": "Packaging portrait image & provisioning LTX-2.3 3D compute engine...",
        "script": script_text,
        "voice": voice,
        "slug": kernel_id,
        "mode": "premium",
        "created_at": time.time(),
        "logs": [f"[{time.strftime('%H:%M:%S')}] Premium LTX-2.3 Job initialized."]
    }
    save_jobs(jobs)
    
    ib64 = ""
    isize = os.path.getsize(image_path)
    if isize <= 500 * 1024 and not hf_repo:
        append_log(job_id, f"Input image ({isize//1024} KB) embedded into script.")
        with open(image_path, "rb") as vf:
            ib64 = base64.b64encode(vf.read()).decode("ascii")
    else:
        append_log(job_id, f"Input image ({isize//1024} KB) will be fetched via dataset URL.")
        
    if hf_repo and hf_token:
        append_log(job_id, f"Uploading source portrait to Hugging Face Dataset {hf_repo}...")
        upload_to_hf_dataset(image_path, hf_repo, f"inputs/{job_id}.png", hf_token)

    script_content = PREMIUM_KERNEL_TEMPLATE.replace("___SCRIPT_TEXT___", repr(script_text)).replace("___VOICE___", repr(voice)).replace("___IMAGE_B64___", repr(ib64)).replace("___HF_REPO___", repr(hf_repo)).replace("___JOB_ID___", repr(job_id))
    with open(os.path.join(staging, "run_epicsync.py"), "w", encoding="utf-8") as f:
        f.write(script_content)
        
    meta = {
        "id": kernel_id,
        "title": f"EpicSync Premium {job_id.split('_')[-1]}",
        "code_file": "run_epicsync.py",
        "language": "python",
        "kernel_type": "script",
        "is_private": True,
        "enable_gpu": True,
        "enable_tpu": False,
        "enable_internet": True,
        "keywords": ["gpu", "diffusion", "ltx"],
        "dataset_sources": [
            "guitammelbader/wan2gp-models",
            "canodian/pl-ltx-2-3-spatial-upscaler-x2-1-0-safetensors"
        ],
        "competition_sources": [],
        "kernel_sources": [],
        "model_sources": [],
        "machine_shape": "NvidiaTeslaT4"
    }
    with open(os.path.join(staging, "kernel-metadata.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)
        
    env = setup_kaggle_auth(kaggle_user, kaggle_key)
    append_log(job_id, f"Pushing Premium kernel {kernel_id} to Kaggle with mounted LTX datasets...")
    
    res = subprocess.run(f"kaggle kernels push -p {staging}", shell=True, capture_output=True, text=True, env=env)
    if res.returncode != 0:
        append_log(job_id, f"ERROR pushing kernel: {res.stderr or res.stdout}")
        jobs = load_jobs()
        jobs[job_id]["status"] = "FAILED"
        save_jobs(jobs)
    else:
        background_tasks.add_task(monitor_job, job_id, kernel_id, env, hf_repo, hf_token)
        
    return {"job_id": job_id, "status": "STAGING"}

@app.get("/api/jobs")
def get_jobs():
    return load_jobs()

@app.post("/api/cancel/{job_id}")
def cancel_job(job_id: str, kaggle_user: str = Form("ikechukwuebiringa1"), kaggle_key: str = Form("KGAT_fc473ab2c166567756eac24217d1fbd2")):
    jobs = load_jobs()
    if job_id not in jobs:
        raise HTTPException(status_code=404, detail="Job not found")
    slug = jobs[job_id].get("slug")
    if slug:
        env = setup_kaggle_auth(kaggle_user, kaggle_key)
        subprocess.run(f"kaggle kernels cancel {slug}", shell=True, env=env)
    jobs[job_id]["status"] = "CANCELLED"
    jobs[job_id]["progress"] = 0
    jobs[job_id]["step_text"] = "Task cancelled by user."
    append_log(job_id, "Job explicitly cancelled by user.")
    save_jobs(jobs)
    return {"status": "CANCELLED"}

@app.post("/api/clear_logs")
def clear_logs():
    jobs = load_jobs()
    # Keep successful runs or clear all logs per user preference
    jobs = {k: v for k, v in jobs.items() if v.get("status") == "RUNNING"}
    save_jobs(jobs)
    return {"status": "CLEARED"}

@app.get("/api/video/{job_id}")
def get_video(job_id: str):
    path = os.path.join(OUTPUTS_DIR, f"{job_id}.mp4")
    if os.path.exists(path):
        return FileResponse(path, media_type="video/mp4")
    raise HTTPException(status_code=404, detail="Video file not found")

@app.get("/api/download/{job_id}")
def download_video(job_id: str):
    path = os.path.join(OUTPUTS_DIR, f"{job_id}.mp4")
    if os.path.exists(path):
        return FileResponse(path, media_type="video/mp4", filename=f"EpicSync_{job_id}.mp4")
    raise HTTPException(status_code=404, detail="Video file not found")

app.mount("/", StaticFiles(directory="static", html=True), name="static")