diff --git a/html/locale_en.json b/html/locale_en.json index d387c860a..049cd9067 100644 --- a/html/locale_en.json +++ b/html/locale_en.json @@ -830,7 +830,7 @@ {"id":"","label":"LTX enable refine","localized":"","hint":"","ui":"video"}, {"id":"","label":"LTX refine strength","localized":"","hint":"","ui":"video"}, {"id":"","label":"LTX decode timestep","localized":"","hint":"","ui":"video"}, - {"id":"","label":"LTX enable audio","localized":"","hint":"","ui":"video"}, + {"id":"","label":"LTX save audio","localized":"","hint":"LTX-2 audio-capable models always generate audio from the same prompt as video; this toggle controls whether the audio track is included in the saved video file","ui":"video"}, {"id":"","label":"Loop","localized":"","hint":"","ui":"extras"}, {"id":"","label":"Local directory name","localized":"","hint":"Directory where to install extension, leave blank for default","ui":"component-8746"}, {"id":"","label":"Libs","localized":"","hint":"","ui":"component-8779"}, diff --git a/modules/ltx/ltx_process.py b/modules/ltx/ltx_process.py index 28c58b2dd..e29c5fd79 100644 --- a/modules/ltx/ltx_process.py +++ b/modules/ltx/ltx_process.py @@ -62,10 +62,12 @@ def _canonical_stage2_kwargs() -> dict: } -def _latent_pass(caps, prompt, negative, width, height, frames, steps, guidance_scale, mp4_fps, conditions, image_cond_noise_scale, seed, image=None): +def _latent_pass(caps, prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask, width, height, frames, steps, guidance_scale, mp4_fps, conditions, image_cond_noise_scale, seed, image=None): base_args = { - 'prompt': prompt, - 'negative_prompt': negative, + 'prompt_embeds': prompt_embeds, + 'prompt_attention_mask': prompt_attention_mask, + 'negative_prompt_embeds': negative_prompt_embeds, + 'negative_prompt_attention_mask': negative_prompt_attention_mask, 'width': get_bucket(width), 'height': get_bucket(height), 'num_frames': get_frames(frames), @@ -93,13 +95,8 @@ def _latent_pass(caps, prompt, negative, width, height, frames, steps, guidance_ base_args['use_cross_timestep'] = True log.debug(f'Video: cls={shared.sd_model.__class__.__name__} op=latent_pass args_keys={list(base_args.keys())}') result = shared.sd_model(**base_args) - # video latents strip the batch dim; audio latents keep it so LTX2Pipeline.prepare_audio_latents - # can rewrap them when re-entered as ndim==4 at Stage 2. latents = result.frames[0] if hasattr(result, 'frames') else None - audio_latents = None - if hasattr(result, 'audio') and result.audio is not None: - audio_latents = result.audio - return latents, audio_latents + return latents def run_ltx(task_id, @@ -269,6 +266,7 @@ def run_ltx(task_id, vae_type='Default', vae_tile_frames=16, ) + processing.fix_seed(p) p.scripts = None p.script_args = None p.do_not_save_grid = True @@ -318,7 +316,6 @@ def run_ltx(task_id, yield None, 'LTX: Generate in progress...' audio = None - stage1_audio_latents = None pixels = None frames_out = None needs_latent_path = upsample_enable or refine_enable @@ -327,10 +324,28 @@ def run_ltx(task_id, if needs_latent_path: prompt_final, negative_final, networks = get_prompts(prompt, negative, styles) extra_networks.activate(p, networks) - latents, stage1_audio_latents = _latent_pass( - caps=caps, + # Encode once and reuse across stages; encode_prompt short-circuits when + # embeds are passed to __call__. CPU park keeps them off GPU between stages. + prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask = shared.sd_model.encode_prompt( prompt=prompt_final, - negative=negative_final, + negative_prompt=negative_final, + do_classifier_free_guidance=True, + device=devices.device, + ) + prompt_embeds = prompt_embeds.cpu() + prompt_attention_mask = prompt_attention_mask.cpu() if prompt_attention_mask is not None else None + negative_prompt_embeds = negative_prompt_embeds.cpu() if negative_prompt_embeds is not None else None + negative_prompt_attention_mask = negative_prompt_attention_mask.cpu() if negative_prompt_attention_mask is not None else None + # encode_prompt outside pipe.__call__ bypasses the post-forward offload hook; + # re-anchor so the text encoder doesn't stay pinned through Stage 1 forward. + shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, force=True, silent=True) + devices.torch_gc(force=True, reason='ltx:encode') + latents = _latent_pass( + caps=caps, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_embeds=negative_prompt_embeds, + negative_prompt_attention_mask=negative_prompt_attention_mask, width=base_w, height=base_h, frames=frames, @@ -339,7 +354,7 @@ def run_ltx(task_id, mp4_fps=mp4_fps, conditions=conditions, image_cond_noise_scale=image_cond_noise_scale if caps.supports_image_cond_noise_scale else None, - seed=int(seed) if seed is not None else -1, + seed=p.seed, image=p.task_args.get('image'), ) else: @@ -348,8 +363,11 @@ def run_ltx(task_id, yield from abort('Video: process_images returned no frames', ok=True, p=p) return pixels = processed.images - if getattr(processed, 'audio', None) is not None: - audio = processed.audio + raw_audio = getattr(processed, 'audio', None) + if raw_audio is not None: + # Strip batch dim from (B, 2, N); write_audio expects (2, N) for the + # transpose-to-interleaved path used by AAC s16. + audio = raw_audio[0].float().cpu() if raw_audio.ndim == 3 else raw_audio.float().cpu() latents = None except AssertionError as e: yield from abort(e, ok=True, p=p) @@ -388,7 +406,7 @@ def run_ltx(task_id, up_args = { 'width': final_w, 'height': final_h, - 'generator': get_generator(int(seed) if seed is not None else -1), + 'generator': get_generator(p.seed), 'output_type': 'latent', } if latents.ndim == 4: @@ -409,7 +427,7 @@ def run_ltx(task_id, 'height': final_h, 'num_frames': get_frames(frames), 'latents_normalized': False, - 'generator': get_generator(int(seed) if seed is not None else -1), + 'generator': get_generator(p.seed), 'output_type': 'latent', } if latents.ndim == 4: @@ -436,13 +454,15 @@ def run_ltx(task_id, # Refine is terminal: let the pipe decode internally so the final VAE pass runs inside # the same offload/cudnn context as a normal generation (matches Generic Video tab). refine_args = { - 'prompt': prompt_final, - 'negative_prompt': negative_final, + 'prompt_embeds': prompt_embeds, + 'prompt_attention_mask': prompt_attention_mask, + 'negative_prompt_embeds': negative_prompt_embeds, + 'negative_prompt_attention_mask': negative_prompt_attention_mask, 'width': final_w, 'height': final_h, 'num_frames': get_frames(frames), 'num_inference_steps': steps, - 'generator': get_generator(int(seed) if seed is not None else -1), + 'generator': get_generator(p.seed), 'callback_on_step_end': diffusers_callback, 'output_type': 'pil', } @@ -457,14 +477,14 @@ def run_ltx(task_id, # Thread Stage-1 I2V init image through Stage 2 so first-frame identity survives refine. if caps.is_i2v and caps.repo_cls_name in ('LTXImageToVideoPipeline', 'LTX2ImageToVideoPipeline') and p.task_args.get('image') is not None: refine_args['image'] = p.task_args['image'] - # Thread Stage-1 audio latents into Stage 2 on 2.x. The video branch cross-attends - # audio every layer; letting prepare_audio_latents fall back to fresh noise biases - # the video branch off-distribution (desaturated output on distilled 2.x). - if caps.family == '2.x': - if stage1_audio_latents is not None: - refine_args['audio_latents'] = stage1_audio_latents.to(device=devices.device) - if caps.use_cross_timestep: - refine_args['use_cross_timestep'] = True + if caps.family == '2.x' and caps.use_cross_timestep: + refine_args['use_cross_timestep'] = True + # output_type='latent' skips the post-loop audio_vae + vocoder pass when audio + # is unwanted; per-step audio cross-attention still runs for video conditioning. + # Internal video decode is also skipped; vae_decode below picks it up. + want_audio = caps.supports_audio and audio_enable + if not want_audio: + refine_args['output_type'] = 'latent' saved_scheduler_stage2 = None try: @@ -502,10 +522,14 @@ def run_ltx(task_id, yield None, 'LTX: Refine in progress...' try: result = shared.sd_model(latents=latents, **refine_args) - pixels = result.frames[0] if hasattr(result, 'frames') else None - if hasattr(result, 'audio') and result.audio is not None: - audio = result.audio[0].float().cpu() - latents = None + out = result.frames[0] if hasattr(result, 'frames') else None + if want_audio: + pixels = out + if hasattr(result, 'audio') and result.audio is not None: + audio = result.audio[0].float().cpu() + latents = None + else: + latents = out except AssertionError as e: yield from abort(e, ok=True, p=p) return @@ -533,14 +557,15 @@ def run_ltx(task_id, extra_networks.deactivate(p) if needs_latent_path and latents is not None: - # Only reached on upsample-without-refine; refine decodes through the pipe and nulls latents. + # Decode any path that leaves latents intact: upsample-without-refine, or + # refine with output_type='latent' (audio_enable=False). shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae'], force=True, silent=True) devices.torch_gc(force=True, reason='ltx:vae') yield None, 'LTX: VAE decode in progress...' try: if torch.is_tensor(latents): # 0.9.x returns raw latents with output_type='latent'; 2.x pre-denormalizes. - frames_out = vae_decode(latents, decode_timestep if caps.supports_decode_timestep else 0.0, int(seed) if seed is not None else -1, denormalize=caps.family == '0.9') + frames_out = vae_decode(latents, decode_timestep if caps.supports_decode_timestep else 0.0, p.seed, denormalize=caps.family == '0.9') else: frames_out = latents except AssertionError as e: diff --git a/modules/ltx/ltx_ui.py b/modules/ltx/ltx_ui.py index 896a89cf7..7385d5351 100644 --- a/modules/ltx/ltx_ui.py +++ b/modules/ltx/ltx_ui.py @@ -109,7 +109,7 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4 audio_accordion = gr.Accordion(open=False, label="Audio", elem_id='ltx_audio_accordion', visible=False) with audio_accordion: with gr.Row(): - audio_enable = gr.Checkbox(label='LTX enable audio', value=False, elem_id="ltx_audio_enable") + audio_enable = gr.Checkbox(label='LTX save audio', value=True, elem_id="ltx_audio_enable") with gr.Column(elem_id='ltx-output-column', scale=2) as _column_output: with gr.Row(): diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 92427c836..562d31fe4 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -452,17 +452,32 @@ def process_refine(p: processing.StableDiffusionProcessing, output): return output +class AudioFrameList(list): + """list subclass with an audio attribute. Plain lists can't hold attributes, + so process_decode uses this when the pipeline output includes audio.""" + audio = None + + +def attach_audio(results, audio): + if audio is None: + return results + wrapped = AudioFrameList(results if isinstance(results, list) else list(results)) + wrapped.audio = audio + return wrapped + + def process_decode(p: processing.StableDiffusionProcessing, output): shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model, exclude=['vae']) if output is not None: if hasattr(output, 'bytes') and output.bytes is not None: log.debug(f'Generated: bytes={len(output.bytes)}') return output + audio = getattr(output, 'audio', None) if not hasattr(output, 'images') and hasattr(output, 'frames'): log.debug(f'Generated: frames={len(output.frames[0])}') output.images = output.frames[0] if output.images is not None and len(output.images) > 0 and isinstance(output.images[0], Image.Image): - return output.images + return attach_audio(output.images, audio) model = shared.sd_model if not is_refiner_enabled(p) else shared.sd_refiner if not hasattr(model, 'vae'): if hasattr(model, 'pipe') and hasattr(model.pipe, 'vae'): @@ -506,8 +521,9 @@ def process_decode(p: processing.StableDiffusionProcessing, output): results = [] else: log.warning('Processing: no results') + audio = None results = [] - return results + return attach_audio(results, audio) def update_pipeline(sd_model, p: processing.StableDiffusionProcessing): diff --git a/modules/video_models/video_save.py b/modules/video_models/video_save.py index db01c605b..f63ae5d3b 100644 --- a/modules/video_models/video_save.py +++ b/modules/video_models/video_save.py @@ -80,22 +80,25 @@ def numpy_to_tensor(images): return tensor +def add_audio_stream(container, audio_sample_rate: int): + # Must be registered before the first container.mux(); avformat_write_header runs there + # and freezes the stream set, after which new streams have time_base=0/0. + audio_stream = container.add_stream("aac", rate=audio_sample_rate) + audio_stream.codec_context.sample_rate = audio_sample_rate + audio_stream.codec_context.layout = "stereo" + audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate) + log.debug(f'Audio: codec={audio_stream.codec_context.name} rate={audio_stream.codec_context.sample_rate} layout={audio_stream.codec_context.layout} base={audio_stream.codec_context.time_base}') + return audio_stream + + def write_audio( container, + audio_stream, samples: torch.Tensor, audio_sample_rate: int, ) -> None: av = check_av() - # create stream - audio_options = { 'time_base': f'1/{audio_sample_rate}' } - audio_stream = container.add_stream("aac", rate=audio_sample_rate, options=audio_options) - audio_stream.codec_context.sample_rate = audio_sample_rate - audio_stream.codec_context.layout = "stereo" audio_stream.codec_context.format = "fltp" - audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate) - # audio_stream.time_base = audio_stream.codec_context.time_base # TODO audio set time-base - log.debug(f'Audio: codec={audio_stream.codec_context.name} rate={audio_stream.codec_context.sample_rate} layout={audio_stream.codec_context.layout} format={audio_stream.codec_context.format} base={audio_stream.codec_context.time_base}') - # init input samples if samples.ndim == 1: samples = samples[:, None] if samples.shape[1] != 2 and samples.shape[0] == 2: @@ -111,13 +114,11 @@ def write_audio( layout="stereo", ) audio_frames.sample_rate = audio_sample_rate - # init resampler audio_resampler = av.audio.resampler.AudioResampler( format=audio_stream.codec_context.format, layout=audio_stream.codec_context.layout, rate=audio_stream.codec_context.sample_rate, ) - # resample pts = 0 for resampled in audio_resampler.resample(audio_frames): resampled.pts = resampled.pts or 0 @@ -126,7 +127,6 @@ def write_audio( for packet in packets: container.mux(packet) pts += resampled.samples - # flush audio encoder for packet in audio_stream.encode(): container.mux(packet) @@ -177,6 +177,7 @@ def atomic_save_video( stream.width = video_array.shape[2] stream.height = video_array.shape[1] stream.pix_fmt = pix_fmt + audio_stream = add_audio_stream(container, aac) if audio is not None else None for img in video_array: frame = av.VideoFrame.from_ndarray(img, format="rgb24") for packet in stream.encode_lazy(frame): @@ -185,9 +186,9 @@ def atomic_save_video( pbar.update(task, advance=1) for packet in stream.encode(): # flush container.mux(packet) - if audio is not None: + if audio_stream is not None: try: - write_audio(container, audio, aac) + write_audio(container, audio_stream, audio, aac) except Exception as e: log.error(f'Video audio encoding: {e}') errors.display(e, 'Audio')