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