durgappc/infinitetalk
0
1import os2import random3import logging4from typing import Any5 6import torch7import gradio as gr8from PIL import Image9 10from utils.model_loader import ModelManager11from utils.gpu_manager import gpu_manager12 13import wan14from wan.utils.utils import cache_image, cache_video, is_video15from wan.utils.multitalk_utils import save_video_ffmpeg16 17 18# =========================19# HOTFIX: Gradio /api_info crash20# =========================21# Fixes: TypeError: argument of type 'bool' is not iterable22# Caused by gradio_client trying to interpret JSON Schema nodes that can be booleans23try:24 import gradio_client.utils as gcu25 26 _old_json_schema_to_python_type = gcu._json_schema_to_python_type27 28 def _json_schema_to_python_type_patched(schema: Any, defs=None):29 if isinstance(schema, bool):30 return "Any"31 return _old_json_schema_to_python_type(schema, defs)32 33 gcu._json_schema_to_python_type = _json_schema_to_python_type_patched34except Exception as e:35 print("gradio_client patch skipped:", e)36 37 38# =========================39# Logging40# =========================41logging.basicConfig(level=logging.INFO)42logger = logging.getLogger(__name__)43 44 45# =========================46# Globals47# =========================48model_manager: ModelManager | None = None49models_loaded = False50 51 52def initialize_models(progress=gr.Progress()):53 """Download/prepare model assets on first use."""54 global model_manager, models_loaded55 56 if models_loaded:57 return58 59 try:60 progress(0.1, desc="Initializing model manager...")61 model_manager = ModelManager()62 63 progress(0.3, desc="Downloading models (first time only)...")64 65 # Pre-download assets (actual heavy loading happens on first inference)66 model_manager.get_wan_model_path()67 model_manager.get_infinitetalk_weights_path()68 model_manager.get_wav2vec_model_path()69 70 models_loaded = True71 progress(1.0, desc="Models ready!")72 logger.info("Models initialized successfully")73 74 except Exception as e:75 logger.exception("Error initializing models")76 raise gr.Error(f"Failed to initialize models: {str(e)}")77 78 79def _set_seed(seed: int) -> int:80 """Set deterministic seeds and return the final seed used."""81 if seed == -1:82 seed = random.randint(0, 99_999_999)83 84 torch.manual_seed(seed)85 if torch.cuda.is_available():86 torch.cuda.manual_seed(seed)87 88 return seed89 90 91def generate_video(92 image_or_video,93 audio_file,94 resolution="480p",95 steps=40,96 audio_guide_scale=3.0,97 seed=-1,98 progress=gr.Progress(),99):100 """101 Generate a talking video from an image OR dub an existing video.102 103 Note: This is a simplified pipeline example. Your real pipeline may use104 wan_pipeline + diffusion steps etc. This version just stitches frames + audio.105 """106 try:107 if not torch.cuda.is_available():108 raise gr.Error("⚠️ GPU not available. This Space requires GPU hardware to generate videos.")109 110 # Ensure models are prepared111 if not models_loaded:112 initialize_models(progress)113 114 progress(0.1, desc="Processing audio...")115 116 progress(0.2, desc="Loading models...")117 # Load models (kept for parity with your structure)118 size = f"infinitetalk-{resolution.replace('p', '')}"119 wan_pipeline = model_manager.load_wan_model(size=size, device="cuda") # noqa: F841120 121 progress(0.3, desc="Processing input...")122 123 # Decide whether the input is a video or image124 if is_video(image_or_video):125 logger.info("Processing video dubbing input...")126 input_frames = cache_video(image_or_video)127 else:128 logger.info("Processing image-to-video input...")129 input_image = Image.open(image_or_video).convert("RGB")130 input_frames = [input_image]131 132 progress(0.4, desc="Generating video...")133 134 seed = _set_seed(int(seed))135 output_path = f"/tmp/output_{seed}.mp4"136 137 # Simplified output save (frames + audio)138 save_video_ffmpeg(139 input_frames,140 output_path,141 audio_file,142 high_quality_save=False,143 )144 145 progress(1.0, desc="Complete!")146 return output_path147 148 except Exception as e:149 logger.exception("Error generating video")150 gpu_manager.cleanup()151 raise gr.Error(f"Generation failed: {str(e)}")152 153 154def create_interface():155 """Create Gradio UI."""156 with gr.Blocks(title="InfiniteTalk - Talking Video Generator") as demo:157 gr.Markdown(158 """159# 🎬 InfiniteTalk - Talking Video Generator160 161Generate realistic talking head videos with accurate lip-sync from images or dub existing videos with new audio!162 163**Note**: First generation may take a few minutes while models download. Subsequent generations are faster.164"""165 )166 167 with gr.Tabs():168 # Tab 1: Image-to-Video169 with gr.Tab("📸 Image-to-Video"):170 gr.Markdown("Transform a static portrait into a talking video")171 172 with gr.Row():173 with gr.Column():174 image_input = gr.Image(175 type="filepath",176 label="Upload Portrait Image (clear face visibility recommended)",177 )178 audio_input = gr.Audio(179 type="filepath",180 label="Upload Audio (MP3, WAV, or FLAC)",181 )182 183 with gr.Accordion("Advanced Settings", open=False):184 resolution = gr.Radio(185 choices=["480p", "720p"],186 value="480p",187 label="Resolution (480p faster, 720p higher quality)",188 )189 steps = gr.Slider(190 minimum=20,191 maximum=50,192 value=40,193 step=1,194 label="Diffusion Steps (more = higher quality but slower)",195 )196 audio_scale = gr.Slider(197 minimum=1.0,198 maximum=5.0,199 value=3.0,200 step=0.5,201 label="Audio Guide Scale (2–4 recommended)",202 )203 seed = gr.Number(value=-1, label="Seed (-1 for random)")204 205 generate_btn = gr.Button("🎬 Generate Video", variant="primary", size="lg")206 207 with gr.Column():208 output_video = gr.Video(label="Generated Video")209 gr.Markdown("**💡 Tip**: Use a high-quality portrait image with clear facial features.")210 211 generate_btn.click(212 fn=generate_video,213 inputs=[image_input, audio_input, resolution, steps, audio_scale, seed],214 outputs=output_video,215 )216 217 # Tab 2: Video Dubbing218 with gr.Tab("🎥 Video Dubbing"):219 gr.Markdown("Dub an existing video with new audio while maintaining natural movements")220 221 with gr.Row():222 with gr.Column():223 video_input = gr.Video(label="Upload Video (with visible face)")224 audio_input_v2v = gr.Audio(225 type="filepath",226 label="Upload New Audio (MP3, WAV, or FLAC)",227 )228 229 with gr.Accordion("Advanced Settings", open=False):230 resolution_v2v = gr.Radio(231 choices=["480p", "720p"],232 value="480p",233 label="Resolution",234 )235 steps_v2v = gr.Slider(236 minimum=20,237 maximum=50,238 value=40,239 step=1,240 label="Diffusion Steps",241 )242 audio_scale_v2v = gr.Slider(243 minimum=1.0,244 maximum=5.0,245 value=3.0,246 step=0.5,247 label="Audio Guide Scale",248 )249 seed_v2v = gr.Number(value=-1, label="Seed")250 251 generate_btn_v2v = gr.Button("🎬 Generate Dubbed Video", variant="primary", size="lg")252 253 with gr.Column():254 output_video_v2v = gr.Video(label="Dubbed Video")255 gr.Markdown("**💡 Tip**: Use a video with consistent face visibility.")256 257 generate_btn_v2v.click(258 fn=generate_video,259 inputs=[video_input, audio_input_v2v, resolution_v2v, steps_v2v, audio_scale_v2v, seed_v2v],260 outputs=output_video_v2v,261 )262 263 gr.Markdown(264 """265---266### About267Powered by InfiniteTalk (Apache 2.0)268 269⚠️ **Note**: This Space requires GPU hardware to generate videos.270"""271 )272 273 return demo274 275 276if __name__ == "__main__":277 demo = create_interface()278 demo.launch()279