| import os |
| import tempfile |
| import torch |
| import gradio as gr |
| import spaces |
| from huggingface_hub import hf_hub_download |
| from safetensors.torch import load_file |
| from diffusers import LTXPipeline, LTXVideoTransformer3DModel |
| from diffusers.utils import export_to_video |
|
|
| BASE_MODEL = "Lightricks/LTX-Video" |
| FINE_TUNED_MODEL = "DollasAndSpence/LTX2.3-ICEdit-Insight" |
| WEIGHT_NAME = "ltx-2.3-edit-insight-dev-fp8.safetensors" |
|
|
| HF_TOKEN = os.environ.get("HF_TOKEN") |
|
|
| print("Downloading checkpoint weights...") |
| ckpt_path = hf_hub_download( |
| repo_id=FINE_TUNED_MODEL, |
| filename=WEIGHT_NAME, |
| token=HF_TOKEN, |
| ) |
|
|
| print("Configuring transformer model...") |
| transformer_config = dict( |
| LTXVideoTransformer3DModel.load_config( |
| BASE_MODEL, |
| subfolder="transformer", |
| token=HF_TOKEN, |
| ) |
| ) |
| |
| transformer_config["attention_head_dim"] = 128 |
| transformer_config["caption_channels"] = 4096 |
|
|
| |
| transformer = LTXVideoTransformer3DModel.from_config(transformer_config) |
|
|
| print("Loading state dict with shape-mismatch handling...") |
| state_dict = load_file(ckpt_path) |
|
|
| cleaned_state_dict = {} |
| for k, v in state_dict.items(): |
| new_key = k.replace("model.diffusion_model.", "").replace("diffusion_model.", "") |
| cleaned_state_dict[new_key] = v |
|
|
| model_state_dict = transformer.state_dict() |
| loaded_state_dict = {} |
|
|
| for name, param in model_state_dict.items(): |
| if name in cleaned_state_dict: |
| ckpt_param = cleaned_state_dict[name] |
| if ckpt_param.shape == param.shape: |
| loaded_state_dict[name] = ckpt_param |
| else: |
| |
| min_dim0 = min(ckpt_param.shape[0], param.shape[0]) |
| if len(param.shape) == 2: |
| min_dim1 = min(ckpt_param.shape[1], param.shape[1]) |
| new_tensor = param.clone() |
| new_tensor[:min_dim0, :min_dim1] = ckpt_param[:min_dim0, :min_dim1] |
| loaded_state_dict[name] = new_tensor |
| else: |
| new_tensor = param.clone() |
| new_tensor[:min_dim0] = ckpt_param[:min_dim0] |
| loaded_state_dict[name] = new_tensor |
| else: |
| |
| loaded_state_dict[name] = param |
|
|
| transformer.load_state_dict(loaded_state_dict, strict=False) |
| print("Transformer weights loaded successfully.") |
|
|
| transformer = transformer.to(dtype=torch.bfloat16) |
|
|
| print("Loading LTXPipeline...") |
| pipe = LTXPipeline.from_pretrained( |
| BASE_MODEL, |
| transformer=transformer, |
| torch_dtype=torch.bfloat16, |
| token=HF_TOKEN, |
| ) |
|
|
| @spaces.GPU(duration=120) |
| def process_video(prompt, input_video_path, inference_steps, guidance_scale): |
| if not input_video_path: |
| raise gr.Error("Please upload an input video.") |
|
|
| pipe.to("cuda") |
|
|
| kwargs = { |
| "prompt": prompt, |
| "video": input_video_path, |
| "num_inference_steps": inference_steps, |
| "guidance_scale": guidance_scale, |
| "output_type": "np", |
| } |
|
|
| try: |
| output = pipe(**kwargs) |
| except Exception as e: |
| raise gr.Error(f"Inference failed: {str(e)}") |
|
|
| out_file = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name |
| export_to_video(output.frames[0], out_file, fps=24) |
|
|
| return out_file |
|
|
| with gr.Blocks(title="LTX2.3 ICEdit-Insight Video Restoration") as demo: |
| gr.Markdown("# 🎥 LTX2.3 ICEdit-Insight Video Restoration") |
| gr.Markdown("Upload an input video and provide an instruction prompt to restore or edit it.") |
| |
| with gr.Row(): |
| with gr.Column(): |
| input_vid = gr.Video(label="Input Video") |
| prompt = gr.Textbox( |
| label="Instruction Prompt", |
| placeholder="e.g., Clean up AI inpainting glitches, remove artifacts...", |
| ) |
| steps = gr.Slider(minimum=1, maximum=50, value=30, step=1, label="Inference Steps") |
| cfg = gr.Slider(minimum=1.0, maximum=15.0, value=4.0, step=0.1, label="Guidance Scale") |
| run_btn = gr.Button("Generate Fix", variant="primary") |
|
|
| with gr.Column(): |
| output_vid = gr.Video(label="Restored Output Video") |
|
|
| run_btn.click( |
| fn=process_video, |
| inputs=[prompt, input_vid, steps, cfg], |
| outputs=output_vid, |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
