mirror of https://github.com/vladmandic/automatic
258 lines
9.8 KiB
Python
258 lines
9.8 KiB
Python
"""
|
|
- 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"<div class='performance'><p>{summary} {memory}</p></div>"
|