From 00c9a2bb753332dcb24246919fefc7cb8c742a5b Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 4 Sep 2024 15:12:46 -0400 Subject: [PATCH] add cogvideox txt2video --- modules/gr_tempdir.py | 2 +- modules/images.py | 6 +- modules/modeldata.py | 2 + modules/processing.py | 6 +- modules/progress.py | 1 - modules/rife/__init__.py | 2 +- modules/sd_models.py | 4 +- scripts/cogvideo.py | 145 ++++++++++++++++++++++++++++++++++++--- scripts/face_details.py | 2 +- 9 files changed, 149 insertions(+), 21 deletions(-) diff --git a/modules/gr_tempdir.py b/modules/gr_tempdir.py index 90a8f7376..0ee15b314 100644 --- a/modules/gr_tempdir.py +++ b/modules/gr_tempdir.py @@ -70,7 +70,7 @@ def pil_to_temp_file(self, img: Image, dir: str, format="png") -> str: # pylint: img.save(name, pnginfo=(metadata if use_metadata else None)) img.already_saved_as = name size = os.path.getsize(name) - shared.log.debug(f'Save temp: image="{name}" resolution={img.width}x{img.height} size={size}') + shared.log.debug(f'Save temp: image="{name}" width={img.width} height={img.height} size={size}') params = ', '.join([f'{k}: {v}' for k, v in img.info.items()]) params = params[12:] if params.startswith('parameters: ') else params with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file: diff --git a/modules/images.py b/modules/images.py index 23fdf86cf..550b08b25 100644 --- a/modules/images.py +++ b/modules/images.py @@ -596,7 +596,7 @@ def atomically_save_image(): shared.log.error(f'Save failed: file="{fn}" format={image_format} args={save_args} {e}') errors.display(e, 'Image save') size = os.path.getsize(fn) if os.path.exists(fn) else 0 - shared.log.info(f'Save: image="{fn}" type={image_format} resolution={image.width}x{image.height} size={size}') + shared.log.info(f'Save: image="{fn}" type={image_format} width={image.width} height={image.height} size={size}') if shared.opts.save_log_fn != '' and len(exifinfo) > 0: fn = os.path.join(paths.data_path, shared.opts.save_log_fn) if not fn.endswith('.json'): @@ -719,7 +719,9 @@ def save_video(p, images, filename = None, video_type: str = 'none', duration: f return None image = images[0] if p is not None: - namegen = FilenameGenerator(p, seed=p.all_seeds[0], prompt=p.all_prompts[0], image=image) + seed = p.all_seeds[0] if getattr(p, 'all_seeds', None) is not None else p.seed + prompt = p.all_prompts[0] if getattr(p, 'all_prompts', None) is not None else p.prompt + namegen = FilenameGenerator(p, seed=seed, prompt=prompt, image=image) else: namegen = FilenameGenerator(None, seed=0, prompt='', image=image) if filename is None and p is not None: diff --git a/modules/modeldata.py b/modules/modeldata.py index aedbf99e2..621fe71d6 100644 --- a/modules/modeldata.py +++ b/modules/modeldata.py @@ -105,6 +105,8 @@ class Shared(sys.modules[__name__].__class__): model_type = 'auraflow' elif "FluxPipeline" in self.sd_model.__class__.__name__ or "FluxControlNetPipeline" in self.sd_model.__class__.__name__: model_type = 'f1' + elif "CogVideoXPipeline" in self.sd_model.__class__.__name__ or "CogVideoXVideoToVideoPipeline": + model_type = 'cogvideox' else: model_type = self.sd_model.__class__.__name__ except Exception: diff --git a/modules/processing.py b/modules/processing.py index 5c7390903..da62de1c9 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -8,6 +8,7 @@ from modules import shared, devices, errors, images, scripts, memstats, lowvram, from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet from modules.processing_class import StableDiffusionProcessing, StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, StableDiffusionProcessingControl # pylint: disable=unused-import from modules.processing_info import create_infotext +from modules.modeldata import model_data from modules import pag @@ -37,7 +38,7 @@ class Processed: self.images = images_list self.prompt = p.prompt self.negative_prompt = p.negative_prompt - self.seed = seed + self.seed = seed if seed != -1 else p.seed self.subseed = subseed self.subseed_strength = p.subseed_strength self.info = info @@ -51,7 +52,7 @@ class Processed: self.batch_size = p.batch_size self.restore_faces = p.restore_faces self.face_restoration_model = shared.opts.face_restoration_model if p.restore_faces else None - self.sd_model_hash = getattr(shared.sd_model, 'sd_model_hash', '') + self.sd_model_hash = getattr(shared.sd_model, 'sd_model_hash', '') if model_data.sd_model is not None else '' self.seed_resize_from_w = p.seed_resize_from_w self.seed_resize_from_h = p.seed_resize_from_h self.denoising_strength = p.denoising_strength @@ -114,7 +115,6 @@ class Processed: return create_infotext(p, self.all_prompts, self.all_seeds, self.all_subseeds, comments=[], position_in_batch=index % self.batch_size, iteration=index // self.batch_size) - def process_images(p: StableDiffusionProcessing) -> Processed: debug(f'Process images: {vars(p)}') if not hasattr(p.sd_model, 'sd_checkpoint_info'): diff --git a/modules/progress.py b/modules/progress.py index aeef195b4..abd6d906d 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -74,7 +74,6 @@ def progressapi(req: ProgressRequest): predicted = elapsed / progress if progress > 0 else None eta = predicted - elapsed if predicted is not None else None # shared.log.debug(f'Progress: step={step_x}:{step_y} batch={batch_x}:{batch_y} current={current} total={total} progress={progress} elapsed={elapsed} eta={eta}') - id_live_preview = req.id_live_preview live_preview = None shared.state.set_current_image() diff --git a/modules/rife/__init__.py b/modules/rife/__init__.py index 7e40735e7..f74f3d984 100644 --- a/modules/rife/__init__.py +++ b/modules/rife/__init__.py @@ -113,5 +113,5 @@ def interpolate(images: list, count: int = 2, scale: float = 1.0, pad: int = 1, while not buffer.empty(): time.sleep(0.1) t1 = time.time() - shared.log.info(f'RIFE interpolate: input={len(images)} frames={len(interpolated)} resolution={w}x{h} interpolate={count} scale={scale} pad={pad} change={change} time={round(t1 - t0, 2)}') + shared.log.info(f'RIFE interpolate: input={len(images)} frames={len(interpolated)} width={w} height={h} interpolate={count} scale={scale} pad={pad} change={change} time={round(t1 - t0, 2)}') return interpolated diff --git a/modules/sd_models.py b/modules/sd_models.py index bc819e53f..f4e7b44a6 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -658,8 +658,8 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False): def copy_diffuser_options(new_pipe, orig_pipe): - new_pipe.sd_checkpoint_info = orig_pipe.sd_checkpoint_info - new_pipe.sd_model_checkpoint = orig_pipe.sd_model_checkpoint + new_pipe.sd_checkpoint_info = getattr(orig_pipe, 'sd_checkpoint_info', None) + new_pipe.sd_model_checkpoint = getattr(orig_pipe, 'sd_model_checkpoint', None) new_pipe.embedding_db = getattr(orig_pipe, 'embedding_db', None) new_pipe.sd_model_hash = getattr(orig_pipe, 'sd_model_hash', None) new_pipe.has_accelerate = getattr(orig_pipe, 'has_accelerate', False) diff --git a/scripts/cogvideo.py b/scripts/cogvideo.py index b8f3ea340..fb5cbdce3 100644 --- a/scripts/cogvideo.py +++ b/scripts/cogvideo.py @@ -1,7 +1,19 @@ +""" +models: https://huggingface.co/THUDM/CogVideoX-2b https://huggingface.co/THUDM/CogVideoX-5b +source: https://github.com/THUDM/CogVideo +quanto: https://gist.github.com/a-r-r-o-w/31be62828b00a9292821b85c1017effa +torchao: https://gist.github.com/a-r-r-o-w/4d9732d17412888c885480c6521a9897 +""" import os +import time import gradio as gr +import torch import diffusers -from modules import scripts, processing, shared, devices, sd_models +from modules import scripts, shared, devices, errors, sd_models, processing +from modules.processing_callbacks import diffusers_callback, set_callbacks_p + + +debug = (os.environ.get('SD_LOAD_DEBUG', None) is not None) or (os.environ.get('SD_PROCESS_DEBUG', None) is not None) class Script(scripts.Script): @@ -24,32 +36,145 @@ class Script(scripts.Script): with gr.Row(): gr.HTML("  CogVideoX
") with gr.Row(): - model = gr.Dropdown(label='Model', choices=['THUDM/CogVideoX-2b', 'THUDM/CogVideoX-5b'], value='THUDM/CogVideoX-2b') + model = gr.Dropdown(label='Model', choices=['None', 'THUDM/CogVideoX-2b', 'THUDM/CogVideoX-5b'], value='THUDM/CogVideoX-2b') sampler = gr.Dropdown(label='Sampler', choices=['DDIM', 'DPM'], value='DDIM') with gr.Row(): - frames = gr.Slider(label='Frames', minimum=1, maximum=64, step=1, value=16) + frames = gr.Slider(label='Frames', minimum=1, maximum=100, step=1, value=49) guidance = gr.Slider(label='Guidance', minimum=0.0, maximum=14.0, step=0.5, value=6.0) with gr.Row(): offload = gr.Dropdown(label='Offload', choices=['none', 'balanced', 'model', 'sequential'], value='balanced') override = gr.Checkbox(label='Override resolution', value=True) with gr.Row(): video_type = gr.Dropdown(label='Video file', choices=['None', 'GIF', 'PNG', 'MP4'], value='None') - duration = gr.Slider(label='Duration', minimum=0.25, maximum=10, step=0.25, value=2, visible=False) + duration = gr.Slider(label='Duration', minimum=0.25, maximum=30, step=0.25, value=8, visible=False) + with gr.Accordion('Optional init image or video', open=False): + with gr.Row(): + image = gr.Image(value=None, label='Image', type='pil', source='upload', width=256, height=256) + video = gr.Video(value=None, label='Video', source='upload', width=256, height=256) with gr.Row(): loop = gr.Checkbox(label='Loop', value=True, visible=False) pad = gr.Slider(label='Pad frames', minimum=0, maximum=24, step=1, value=1, visible=False) interpolate = gr.Slider(label='Interpolate frames', minimum=0, maximum=24, step=1, value=0, visible=False) video_type.change(fn=video_type_change, inputs=[video_type], outputs=[duration, loop, pad, interpolate]) - return [model, sampler, frames, guidance, offload, override, video_type, duration, duration, loop, pad, interpolate] + return [model, sampler, frames, guidance, offload, override, video_type, duration, loop, pad, interpolate, image, video] - def run(self, p: processing.StableDiffusionProcessing, model, sampler, frames, guidance, offload, override, video_type, duration, loop, pad, interpolate): # pylint: disable=arguments-differ, unused-argument - shared.log.debug(f'CogVideoX: model={model} sampler={sampler} frames={frames} guidance={guidance} offload={offload} override={override} video_type={video_type} duration={duration} loop={loop} pad={pad} interpolate={interpolate}') + def load(self, model, txt): + if shared.sd_model_type != 'cogvideox' and model != 'None': + sd_models.unload_model_weights('model') + shared.log.info(f'CogVideoX load: model="{model}"') + try: + shared.sd_model = diffusers.CogVideoXPipeline.from_pretrained(model, torch_dtype=devices.dtype, cache_dir=shared.opts.diffusers_dir) + shared.sd_model.sd_checkpoint_info = sd_models.CheckpointInfo(model) + shared.sd_model.sd_model_checkpoint = model + except Exception as e: + shared.log.error(f'Loading CogVideoX: {e}') + if debug: + errors.display(e, 'CogVideoX') + devices.torch_gc() + if shared.sd_model_type == 'cogvideox' and model != 'None': + shared.sd_model = sd_models.switch_pipe(diffusers.CogVideoXPipeline if txt else diffusers.CogVideoXVideoToVideoPipeline, shared.sd_model) + shared.sd_model.set_progress_bar_config(bar_format='Progress {rate_fmt}{postfix} {bar} {percentage:3.0f}% {n_fmt}/{total_fmt} {elapsed} {remaining} ' + '\x1b[38;5;71m', ncols=80, colour='#327fba') + shared.log.debug(f'CogVideoX load: class="{shared.sd_model.__class__.__name__}"') + if shared.sd_model is not None and model == 'None': + shared.log.info(f'CogVideoX unload: model={model}') + shared.sd_model = None + devices.torch_gc(force=True) + + def offload(self, offload): + if shared.sd_model_type != 'cogvideox': + return + if offload == 'none': + sd_models.move_model(shared.sd_model, devices.device) + shared.log.info(f'CogVideoX: offload={offload}') + if offload == 'balanced': + sd_models.apply_balanced_offload(shared.sd_model) + if offload == 'model': + shared.sd_model.enable_model_cpu_offload() + if offload == 'sequential': + shared.sd_model.enable_model_cpu_offload() + shared.sd_model.enable_sequential_cpu_offload() + shared.sd_model.vae.enable_slicing() + shared.sd_model.vae.enable_tiling() + + def generate(self, p: processing.StableDiffusionProcessing): + if shared.sd_model_type != 'cogvideox': + return [] + shared.log.info(f'CogVideoX: sampler={p.sampler} steps={p.steps} frames={p.frames} width={p.width} height={p.height} seed={p.seed} guidance={p.guidance}') + if p.sampler == 'DDIM': + shared.sd_model.scheduler = diffusers.CogVideoXDDIMScheduler.from_config(shared.sd_model.scheduler.config, timestep_spacing="trailing") + if p.sampler == 'DPM': + shared.sd_model.scheduler = diffusers.CogVideoXDPMScheduler.from_config(shared.sd_model.scheduler.config, timestep_spacing="trailing") + t0 = time.time() + frames = [] + set_callbacks_p(p) + shared.state.job_count = 1 + shared.state.sampling_steps = p.steps - 1 + try: + args = dict( + prompt=p.prompt, + negative_prompt=p.negative_prompt, + height=p.height, + width=p.width, + num_videos_per_prompt=1, + num_inference_steps=p.steps, + guidance_scale=p.guidance, + generator=torch.Generator(device=devices.device).manual_seed(p.seed), + callback_on_step_end=diffusers_callback, + callback_on_step_end_tensor_inputs=['latents'], + ) + if getattr(p, 'image', False): + raise ValueError('CogVideoX: image not supported') + # args['latents'] = [p.image] + elif getattr(p, 'video', False): + raise ValueError('CogVideoX: video not supported') + # args['video'] = p.video + else: + args['num_frames'] = p.frames # only txt2vid has num_frames + if debug: + shared.log.debug(f'CogVideoX args: {args}') + frames = shared.sd_model(**args).frames[0] + except AssertionError as e: + shared.log.info(f'CogVideoX: {e}') + except Exception as e: + shared.log.error(f'CogVideoX: {e}') + if debug: + errors.display(e, 'CogVideoX') + t1 = time.time() + its = (len(frames) * p.steps) / (t1 - t0) + shared.log.info(f'CogVideoX: frames={len(frames)} its={its:.2f} time={t1 - t0:.2f}') + return frames + + # auto-executed by the script-callback + def run(self, p: processing.StableDiffusionProcessing, model, sampler, frames, guidance, offload, override, video_type, duration, loop, pad, interpolate, image, video): # pylint: disable=arguments-differ, unused-argument + shared.state.begin('CogVideoX') + processing.fix_seed(p) p.extra_generation_params['CogVideoX'] = model p.do_not_save_grid = True if 'animatediff' not in p.ops: p.ops.append('cogvideox') + if override: + p.width = 720 + p.height = 480 + p.sampler = sampler + p.guidance = guidance + p.frames = frames + p.use_dynamic_cfg = sampler == 'DPM' + p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles) + p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles) + p.image = image + p.video = video + txt = image is None and video is None + self.load(model, txt) + self.offload(offload) + frames = self.generate(p) + info = 'whatever' + processed = processing.Processed(p, images_list=frames, info=info) + shared.state.end() + return processed - def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, model, sampler, frames, guidance, override_resolution, video_type, duration, loop, pad, interpolate): # pylint: disable=arguments-differ, unused-argument - from modules.images import save_video - if video_type != 'None': + # auto-executed by the script-callback + def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, model, sampler, frames, guidance, offload, override, video_type, duration, loop, pad, interpolate, image, video): # pylint: disable=arguments-differ, unused-argument + if video_type != 'None' and processed is not None and len(processed.images) > 0: + from modules.images import save_video + shared.log.info(f'CogVideoX video: type={video_type} duration={duration} loop={loop} pad={pad} interpolate={interpolate}') save_video(p, filename=None, images=processed.images, video_type=video_type, duration=duration, loop=loop, pad=pad, interpolate=interpolate) diff --git a/scripts/face_details.py b/scripts/face_details.py index 004c775f6..7aca8a528 100644 --- a/scripts/face_details.py +++ b/scripts/face_details.py @@ -161,7 +161,7 @@ class FaceRestorerYolo(FaceRestoration): p.negative_prompt = orig_p.get('all_negative_prompts', [''])[0] report = [{'score': f.score, 'size': f'{f.width}x{f.height}' } for f in faces] - shared.log.debug(f'Face HiRes: faces={report} args={faces[0].args} denoise={p.denoising_strength} blur={p.mask_blur} resolution={p.width}x{p.height} padding={p.inpaint_full_res_padding}') + shared.log.debug(f'Face HiRes: faces={report} args={faces[0].args} denoise={p.denoising_strength} blur={p.mask_blur} width={p.width} height={p.height} padding={p.inpaint_full_res_padding}') mask_all = [] for face in faces: