kaggle_template/run_batch.py

import os
import sys
import json
import random
import tempfile
import glob
import traceback
import shutil
import subprocess
import psutil
import soundfile as sf
import numpy as np
if not hasattr(np, '_no_nep50_warning'):
    def _dummy_nep50(func=None, *args, **kwargs):
        if func is not None and callable(func):
            return func
        return lambda f: f
    np._no_nep50_warning = _dummy_nep50
from PIL import Image

# --- EMBEDDED JOB DATA ---
EMBEDDED_CONFIG = None
EMBEDDED_SCRIPT = None
EMBEDDED_PROMPT = None
EMBEDDED_IMAGE_URL = None
EMBEDDED_IMAGE_EXT = ".png"
# -------------------------

print("=== Starting LTX Flow B Headless Batch Runner ===")
sys.stdout.flush()

# Step 1: Install dependencies
print("Installing dependencies...")
sys.stdout.flush()
subprocess.run(["pip", "install", "-q", "torch==2.3.1", "torchvision==0.18.1", "torchaudio==2.3.1"], check=True)
if not os.path.exists("Wan2GP"):
    subprocess.run(["git", "clone", "https://github.com/deepbeepmeep/Wan2GP.git"], check=True)
subprocess.run(["pip", "install", "--timeout", "120", "--retries", "5", "-q", "-r", "Wan2GP/requirements.txt"], check=True)
subprocess.run(["pip", "install", "--timeout", "120", "--retries", "5", "-q", "mmgp", "gguf", "soundfile", "edge-tts", "huggingface_hub", "transformers==4.48.3"], check=True)

# ---- bootstrap Wan2GP ----
WAN2GP_DIR = os.path.abspath("Wan2GP")
sys.path.insert(0, WAN2GP_DIR)
os.chdir(WAN2GP_DIR)
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True,max_split_size_mb:128,garbage_collection_threshold:0.5"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["CUDA_LAUNCH_BLOCKING"] = "0"

import torch

# --- Monkey-patch for transformers backward compatibility with optimum.quanto ---
try:
    import transformers
    import transformers.generation
    if not hasattr(transformers.generation, "GenerationMixin"):
        from transformers.generation.utils import GenerationMixin
        transformers.generation.GenerationMixin = GenerationMixin
        print("Successfully monkey-patched transformers.generation.GenerationMixin")
except Exception as e:
    print(f"Transformers patch notice: {e}")
# --------------------------------------------------------------------------------

# Step 2: Download Models to /kaggle/tmp and symlink
print("Downloading models...")
sys.stdout.flush()
from huggingface_hub import hf_hub_download

REPO = 'DeepBeepMeep/LTX-2'
MODEL_DIR = 'models'
TMP_DIR = '/kaggle/tmp/models'
os.makedirs(MODEL_DIR, exist_ok=True)
os.makedirs(TMP_DIR, exist_ok=True)

LARGE_FILES = [
    'ltx-2.3-22b-distilled-Q4_K_M_light.gguf',
    'ltx-2.3-22b-distilled-lora-384.safetensors',
    'ltx-2.3-22b_embeddings_connector.safetensors',
    'ltx-2.3-22b_text_embedding_projection.safetensors',
    'ltx-2.3-22b_vae.safetensors',
    'id-lora-celebvhq-ltx2.3.safetensors',
]

for f in LARGE_FILES:
    dest = os.path.join(MODEL_DIR, f)
    if os.path.exists(dest): continue
    print(f'Downloading {f} -> /kaggle/tmp ...')
    hf_hub_download(repo_id=REPO, filename=f, local_dir=TMP_DIR)
    actual = os.path.join(TMP_DIR, f)
    if not os.path.exists(dest):
        os.symlink(actual, dest)

SMALL_FILES = [
    'ltx-2.3-22b_audio_vae.safetensors',
    'ltx-2.3-22b_vocoder.safetensors',
    'ltx-2.3-spatial-upscaler-x2-1.1.safetensors',
    'ltx-2.3-temporal-upscaler-x2-1.0.safetensors',
]

for f in SMALL_FILES:
    dest = os.path.join(MODEL_DIR, f)
    if os.path.exists(dest): continue
    print(f'Downloading {f}...')
    hf_hub_download(repo_id=REPO, filename=f, local_dir=MODEL_DIR)

GEMMA_FOLDER = 'gemma-3-12b-it-qat-q4_0-unquantized'
GEMMA_FILES = [
    'gemma-3-12b-it-qat-q4_0-unquantized_quanto_bf16_int8.safetensors',
    'added_tokens.json', 'chat_template.json', 'config_light.json',
    'generation_config.json', 'preprocessor_config.json', 'processor_config.json',
    'special_tokens_map.json', 'tokenizer.json', 'tokenizer.model', 'tokenizer_config.json',
]

gemma_dest = os.path.join(MODEL_DIR, GEMMA_FOLDER)
gemma_tmp = os.path.join(TMP_DIR, GEMMA_FOLDER)
if not os.path.exists(gemma_dest):
    os.makedirs(gemma_tmp, exist_ok=True)
    for gf in GEMMA_FILES:
        tmp_file = os.path.join(gemma_tmp, gf)
        if os.path.exists(tmp_file): continue
        print(f'Downloading gemma/{gf} -> /kaggle/tmp ...')
        hf_hub_download(repo_id=REPO, filename=f'{GEMMA_FOLDER}/{gf}', local_dir=TMP_DIR)
    os.symlink(gemma_tmp, gemma_dest)

print("All models ready!")
sys.stdout.flush()

# Step 3: Register GGUF & Load Model
def _register_gguf_handler():
    import shared.qtypes.gguf
    print("  [GGUF] Extension handler registered")

def _patch_ltx2_config_loading():
    import models.ltx2.ltx2 as ltx2_mod
    _original = ltx2_mod._load_config_from_checkpoint
    def _patched(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"):
                config = ltx2_mod._normalize_config(metadata.get("config"))
                if config: return config
        except Exception: pass
        if fallback_config_path and os.path.isfile(fallback_config_path):
            try:
                with open(fallback_config_path, "r", encoding="utf-8") as f:
                    return ltx2_mod._normalize_config(json.load(f))
            except Exception: pass
        return {}
    ltx2_mod._load_config_from_checkpoint = _patched

_register_gguf_handler()
_patch_ltx2_config_loading()

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
from shared.utils.audio_video import save_video

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_folder = "models/gemma-3-12b-it-qat-q4_0-unquantized"
gemma_files = sorted(glob.glob(os.path.join(gemma_folder, "*.safetensors")))
quanto_files = [f for f in gemma_files if "quanto" in f]
text_encoder_file = quanto_files[0] if quanto_files else (gemma_files[0] if gemma_files else None)
transformer_path = os.path.join("models", "ltx-2.3-22b-distilled-Q4_K_M_light.gguf")

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,
)

offload.profile(
    pipe,
    profile_no=4,
    quantizeTransformer=False,
    convertWeightsFloatTo=torch.float16,
    budgets={
        "transformer": 6000, "text_encoder": 1500, "video_encoder": 2000,
        "video_decoder": 3000, "audio_encoder": 1000, "audio_decoder": 1000,
        "vocoder": 500, "spatial_upsampler": 1500, "vae": 1000, "*": 1000,
    },
)
offload.shared_state["_attention"] = "sdpa"
print("Model loaded & profile 4 applied!")
sys.stdout.flush()

# Step 4: Helpers
def get_resolution(base_res_str, aspect_ratio_str):
    base_resolutions = {"1080p": 1088, "720p": 704, "540p": 544, "480p": 480}
    ratios = {"16:9 Landscape": 16/9, "4:3 Standard": 4/3, "1:1 Square": 1.0, "3:4 Portrait": 3/4, "9:16 Portrait": 9/16}
    base = base_resolutions.get(base_res_str, 704)
    ratio = ratios.get(aspect_ratio_str, 16/9)
    width = int(base * ratio) if ratio >= 1.0 else base
    height = base if ratio >= 1.0 else int(base / ratio)
    return (width // 32) * 32, (height // 32) * 32

def snap_to_ltx_frames(duration_sec: float, fps: float = 24.0, max_frames: int = 721) -> int:
    raw = duration_sec * fps
    k = max(0, round((raw - 1) / 8))
    frames = 8 * k + 1
    return int(max(49, min(frames, max_frames)))

@torch.inference_mode()
def Video_Generation(prompt, input_image_start, input_audio, seed, resolution_dropdown, aspect_ratio_dropdown, guide_scale=4.0):
    import gc
    gc.collect(); torch.cuda.empty_cache(); torch.cuda.synchronize()
    frame_rate = 24.0
    wav, sr = sf.read(input_audio)
    if wav.ndim > 1: wav = wav.mean(axis=1)
    input_waveform = wav.astype(np.float32)
    audio_duration_sec = len(wav) / sr
    num_frames = snap_to_ltx_frames(audio_duration_sec, frame_rate)
    width, height = get_resolution(resolution_dropdown, aspect_ratio_dropdown)
    
    if seed is None or seed < 0: seed = random.randint(0, 2**32 - 1)
    image_start = Image.open(input_image_start).convert("RGB") if input_image_start else None

    print(f"Generating video clip: {width}x{height}, {num_frames}f ({audio_duration_sec:.1f}s), seed={seed}")
    sys.stdout.flush()

    total_steps = [8]; current_step = [0]; current_pass = [1]
    def cb(step, latent, is_start, override_num_inference_steps=None, pass_no=None, **kwargs):
        if is_start:
            if override_num_inference_steps is not None: total_steps[0] = override_num_inference_steps
            if pass_no is not None: current_pass[0] = pass_no
            current_step[0] = 0
            return
        current_step[0] += 1
        print(f"  [Pass {current_pass[0]}] step {current_step[0]}/{total_steps[0]}")
        sys.stdout.flush()

    gen_kwargs = dict(
        input_prompt=prompt, image_start=image_start, height=height, width=width,
        frame_num=num_frames, fps=frame_rate, seed=seed, callback=cb, VAE_tile_size=256,
        input_video_strength=1.0, denoising_strength=1.0, guide_scale=float(guide_scale),
        sampling_steps=8, guide_phases=2, n_prompt="", video_prompt_type="", audio_prompt_type="2",
        input_waveform=input_waveform, input_waveform_sample_rate=int(sr), audio_scale=1.0
    )
    def set_progress_status(status: str): print(f"  [{status}]..."); sys.stdout.flush()
    gen_kwargs["set_progress_status"] = set_progress_status

    result = ltx2_model.generate(**gen_kwargs)
    if result is None: raise RuntimeError("Generation returned None")

    video_tensor = result.get("x") if isinstance(result, dict) else (result[0] if isinstance(result, tuple) else result)
    audio_data = result.get("audio", input_waveform) if isinstance(result, dict) else (result[1] if isinstance(result, tuple) and len(result)>1 else input_waveform)
    audio_sr = result.get("audio_sampling_rate", int(sr)) if isinstance(result, dict) else (result[2] if isinstance(result, tuple) and len(result)>2 else int(sr))

    out_path = tempfile.mktemp(suffix=".mp4")
    video_for_save = video_tensor.cpu().unsqueeze(0).float() / 127.5 - 1.0
    save_video(tensor=video_for_save, save_file=out_path, fps=frame_rate, normalize=True, value_range=(-1, 1))

    audio_tmp = tempfile.mktemp(suffix=".wav")
    audio_np = audio_data if isinstance(audio_data, np.ndarray) else audio_data.cpu().float().numpy()
    if audio_np.ndim == 2 and audio_np.shape[0] <= 2: audio_np = audio_np.T
    sf.write(audio_tmp, audio_np, int(audio_sr))

    final_path = out_path.replace(".mp4", "_with_audio.mp4")
    subprocess.run(["ffmpeg", "-y", "-i", out_path, "-i", audio_tmp, "-c:v", "copy", "-c:a", "aac", "-b:a", "192k", "-shortest", final_path], check=True, capture_output=True)
    return final_path if os.path.exists(final_path) else out_path

def generate_tts(text, index, voice="en-US-AnaNeural"):
    out_file = f"{index}.wav"
    subprocess.run(["edge-tts", "--text", text, "--voice", voice, "--write-media", out_file], check=True)
    return out_file

# Step 5: Read input config and execute Flow B
input_dir = "/kaggle/working"

if EMBEDDED_CONFIG is not None:
    print("Using embedded self-contained job configuration...")
    cfg = EMBEDDED_CONFIG
    script_text = EMBEDDED_SCRIPT
    prompt_text = EMBEDDED_PROMPT
    if EMBEDDED_IMAGE_URL:
        import urllib.request
        image_file = os.path.join(input_dir, f"reference_image{EMBEDDED_IMAGE_EXT}")
        urllib.request.urlretrieve(EMBEDDED_IMAGE_URL, image_file)
        print(f"Downloaded embedded reference image from {EMBEDDED_IMAGE_URL} to {image_file}")
    else:
        image_file = None
else:
    for d in [os.path.dirname(os.path.abspath(__file__)), "/kaggle/src", "/kaggle/working", "."]:
        if os.path.exists(os.path.join(d, "config.json")):
            input_dir = d
            break
    else:
        input_dir = "/kaggle/working"

    print(f"Loading input configuration from: {input_dir}")
    config_file = os.path.join(input_dir, "config.json")
    script_file = os.path.join(input_dir, "script.txt")
    prompt_file = os.path.join(input_dir, "prompt.txt")

    images = glob.glob(os.path.join(input_dir, "*.*"))
    image_file = next((f for f in images if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')) and not f.endswith('_video.mp4')), None)

    with open(config_file, "r", encoding="utf-8") as f: cfg = json.load(f)
    with open(script_file, "r", encoding="utf-8") as f: script_text = f.read()
    with open(prompt_file, "r", encoding="utf-8") as f: prompt_text = f.read()

paragraphs = [p.strip() for p in script_text.replace('\r', '').split('\n\n') if p.strip()]
if not paragraphs: paragraphs = [p.strip() for p in script_text.replace('\r', '').split('\n') if p.strip()]

generated_videos = []
for i, para in enumerate(paragraphs):
    print(f"\n--- Processing Clip {i+1}/{len(paragraphs)} ---")
    audio_file = generate_tts(para, i+1, voice=cfg.get("voice", "en-US-AnaNeural"))
    vid_out = Video_Generation(
        prompt=prompt_text, input_image_start=image_file, input_audio=audio_file,
        seed=cfg.get("seed", -1), resolution_dropdown=cfg.get("resolution", "720p"),
        aspect_ratio_dropdown=cfg.get("aspect_ratio", "16:9 Landscape"), guide_scale=cfg.get("guide_scale", 4.0)
    )
    final_clip = os.path.join("/kaggle/working", f"{i+1}_video.mp4")
    if os.path.exists(final_clip): os.remove(final_clip)
    shutil.move(vid_out, final_clip)
    generated_videos.append(final_clip)
    print(f"Clip {i+1} saved: {final_clip}")

print("\nStitching final video...")
list_file = "concat_list.txt"
with open(list_file, "w") as f:
    for vid in generated_videos: f.write(f"file '{vid}'\n")

final_output = "/kaggle/working/final_stitched_video.mp4"
if os.path.exists(final_output): os.remove(final_output)
subprocess.run(["ffmpeg", "-y", "-f", "concat", "-safe", "0", "-i", list_file, "-c", "copy", final_output], check=True)

print(f"\nSUCCESS! Final video created at: {final_output}")
sys.stdout.flush()