""" - modernui - teacache and others """ import os import time import threading from modules import shared, errors, timer, memstats, progress, processing, sd_models, sd_samplers, extra_networks from modules.video_models.video_save import save_video from modules.video_models.video_utils import check_av from modules.processing_callbacks import diffusers_callback from modules.ltx.ltx_util import get_bucket, get_frames, load_model, load_upsample, get_conditions, get_generator, get_prompts, vae_decode debug = shared.log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None engine, model = 'LTX Video', 'LTXVideo 0.9.7 13B' upsample_repo_id = "a-r-r-o-w/LTX-Video-0.9.7-Latent-Spatial-Upsampler-diffusers" upsample_pipe = None queue_lock = threading.Lock() def run_ltx(task_id, _ui_state, prompt:str, negative:str, styles:list[str], width:int, height:int, frames:int, steps:int, sampler_index:int, seed:int, upsample_enable:bool, upsample_ratio:float, refine_enable:bool, refine_strength:float, condition_strength: float, condition_image, condition_files, condition_video, condition_video_frames:int, condition_video_skip:int, decode_timestep:float, image_cond_noise_scale:float, mp4_fps:int, mp4_interpolate:int, mp4_codec:str, mp4_ext:str, mp4_opt:str, mp4_video:bool, mp4_frames:bool, mp4_sf:bool, _overrides, ): def abort(e, ok:bool=False, p=None): if ok: shared.log.info(e) else: shared.log.error(f'Video: cls={shared.sd_model.__class__.__name__} op=base {e}') errors.display(e, 'LTX') if p is not None: extra_networks.deactivate(p) shared.state.end() progress.finish_task(task_id) yield None, f'LTX Error: {str(e)}' from diffusers import LTXConditionPipeline # pylint: disable=unused-import check_av() progress.add_task_to_queue(task_id) with queue_lock: progress.start_task(task_id) memstats.reset_stats() timer.process.reset() yield None, 'LTX: Loading...' load_model(engine, model) shared.state.begin('Video', task_id=task_id) shared.state.job_count = 1 p = processing.StableDiffusionProcessingVideo( prompt=prompt, negative_prompt=negative, styles=styles, width=width, height=height, frames=frames, steps=steps, sampler_index=sampler_index, seed=seed, ) p.ops.append('video') conditions = get_conditions( width, height, condition_strength, condition_image, condition_files, condition_video, condition_video_frames, condition_video_skip, ) prompt, negative, networks = get_prompts(prompt, negative, styles) sampler_name = processing.get_sampler_name(sampler_index) sd_samplers.create_sampler(sampler_name, shared.sd_model) shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} op=init prompt="{prompt}" negative="{negative}" styles={styles} networks={networks} sampler={shared.sd_model.scheduler.__class__.__name__}') extra_networks.activate(p, networks) t0 = time.time() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) t1 = time.time() base_args = { "prompt": prompt, "negative_prompt": negative, "width": get_bucket(width), "height": get_bucket(height), "num_frames": get_frames(frames), "num_inference_steps": steps, "image_cond_noise_scale": image_cond_noise_scale, "generator": get_generator(seed), "callback_on_step_end": diffusers_callback, "output_type": "latent", } shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} op=base {base_args}') if len(conditions) > 0: base_args["conditions"] = conditions yield None, 'LTX: Generate in progress...' try: latents = shared.sd_model(**base_args).frames[0] except AssertionError as e: yield from abort(e, ok=True, p=p) return except Exception as e: yield from abort(e, ok=False, p=p) return t2 = time.time() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) t3 = time.time() timer.process.add('offload', t1 - t0) timer.process.add('base', t2 - t1) timer.process.add('offload', t3 - t2) if upsample_enable: t4 = time.time() shared.state.begin('Upsample') global upsample_pipe # pylint: disable=global-statement upsample_pipe = load_upsample(upsample_pipe, upsample_repo_id) upsample_pipe = sd_models.apply_balanced_offload(upsample_pipe) upscale_args = { "width": get_bucket(upsample_ratio * width), "height": get_bucket(upsample_ratio * height), "generator": get_generator(seed), "output_type": "latent", } if latents.ndim == 4: latents = latents.unsqueeze(0) # add batch dimension shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} op=upsample latents={latents.shape} {upscale_args}') yield None, 'LTX: Upsample in progress...' try: upsampled_latents = upsample_pipe(latents=latents, **upscale_args).frames[0] except AssertionError as e: yield from abort(e, ok=True, p=p) return except Exception as e: yield from abort(e, ok=False, p=p) return latents = upsampled_latents t5 = time.time() upsample_pipe = sd_models.apply_balanced_offload(upsample_pipe) t6 = time.time() timer.process.add('upsample', t5 - t4) timer.process.add('offload', t6 - t5) shared.state.end() if refine_enable: t7 = time.time() shared.state.begin('Refine') shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) refine_args = { "prompt": prompt, "negative_prompt": negative, "width": get_bucket(upsample_ratio * width), "height": get_bucket(upsample_ratio * height), "num_frames": get_frames(frames), "denoise_strength": refine_strength, "num_inference_steps": steps, "image_cond_noise_scale": image_cond_noise_scale, "generator": get_generator(seed), "callback_on_step_end": diffusers_callback, "output_type": "latent", } if latents.ndim == 4: latents = latents.unsqueeze(0) # add batch dimension shared.log.debug(f'Video: cls={shared.sd_model.__class__.__name__} op=refine latents={latents.shape} {refine_args}') if len(conditions) > 0: refine_args["conditions"] = conditions yield None, 'LTX: Refine in progress...' try: refined_latents = shared.sd_model(latents=latents, **refine_args).frames[0] except AssertionError as e: yield from abort(e, ok=True, p=p) return except Exception as e: yield from abort(e, ok=False, p=p) return latents = refined_latents t8 = time.time() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) t9 = time.time() timer.process.add('refine', t8 - t7) timer.process.add('offload', t9 - t8) shared.state.end() extra_networks.deactivate(p) yield None, 'LTX: VAE decode in progress...' try: frames = vae_decode(latents, decode_timestep, seed) except AssertionError as e: yield from abort(e, ok=True, p=p) return except Exception as e: yield from abort(e, ok=False, p=p) return t10 = time.time() shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model) t11 = time.time() timer.process.add('offload', t11 - t10) num_frames, video_file = save_video( pixels=frames, mp4_fps=mp4_fps, mp4_codec=mp4_codec, mp4_opt=mp4_opt, mp4_ext=mp4_ext, mp4_sf=mp4_sf, mp4_video=mp4_video, mp4_frames=mp4_frames, mp4_interpolate=mp4_interpolate, metadata={}, ) t_end = time.time() _n, _c, _t, h, w = frames.shape resolution = f'{w}x{h}' if num_frames > 0 else None summary = timer.process.summary(min_time=0.25, total=False).replace('=', ' ') memory = shared.mem_mon.summary() fps = f'{num_frames/(t_end-t0):.2f}' its = f'{(steps)/(t_end-t0):.2f}' shared.state.end() progress.finish_task(task_id) shared.log.info(f'Processed: fn="{video_file}" frames={num_frames} fps={fps} its={its} resolution={resolution} time={t_end-t0:.2f} timers={timer.process.dct()} memory={memstats.memory_stats()}') yield video_file, f'LTX: Generation completed | File {video_file} | Frames {len(frames)} | Resolution {resolution} | f/s {fps} | it/s {its} '+ f"
{summary} {memory}