From db6a52a7fca35cb3106604081f53c2c2a683f941 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 4 Sep 2024 12:32:57 -0400 Subject: [PATCH] update flux --- modules/model_flux.py | 6 ++--- modules/scripts.py | 4 ++-- scripts/cogvideo.py | 55 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 60 insertions(+), 5 deletions(-) create mode 100644 scripts/cogvideo.py diff --git a/modules/model_flux.py b/modules/model_flux.py index 01b093712..e0109cb07 100644 --- a/modules/model_flux.py +++ b/modules/model_flux.py @@ -47,7 +47,7 @@ def load_flux_quanto(checkpoint_info, diffusers_load_config): debug(f'Loading FLUX: quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="transformer"') if not os.path.exists(quantization_map): repo_id = checkpoint_info.name.replace('Diffusers/', '').replace('Diffusers\\', '').replace('models--', '').replace('--', '/') - quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', **diffusers_load_config) + quantization_map = hf_hub_download(repo_id, subfolder='transformer', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) with open(quantization_map, "r", encoding='utf8') as f: quantization_map = json.load(f) state_dict = load_file(os.path.join(repo_path, "transformer", "diffusion_pytorch_model.safetensors")) @@ -72,7 +72,7 @@ def load_flux_quanto(checkpoint_info, diffusers_load_config): debug(f'Loading FLUX: quantization map="{quantization_map}" repo="{checkpoint_info.name}" component="text_encoder_2"') if not os.path.exists(quantization_map): repo_id = checkpoint_info.name.replace('Diffusers/', '').replace('Diffusers\\', '').replace('models--', '').replace('--', '/') - quantization_map = hf_hub_download(repo_id, subfolder='text_encoder_2', filename='quantization_map.json', **diffusers_load_config) + quantization_map = hf_hub_download(repo_id, subfolder='text_encoder_2', filename='quantization_map.json', cache_dir=shared.opts.diffusers_dir) with open(quantization_map, "r", encoding='utf8') as f: quantization_map = json.load(f) with open(os.path.join(repo_path, "text_encoder_2", "config.json"), encoding='utf8') as f: @@ -97,7 +97,7 @@ def load_flux_quanto(checkpoint_info, diffusers_load_config): return transformer, text_encoder_2 -def load_flux_bnb(checkpoint_info, diffusers_load_config, ): # pylint: disable=unused-argument +def load_flux_bnb(checkpoint_info, diffusers_load_config): # pylint: disable=unused-argument transformer, text_encoder_2 = None, None if isinstance(checkpoint_info, str): repo_path = checkpoint_info diff --git a/modules/scripts.py b/modules/scripts.py index 9e92b7f35..464117c5e 100644 --- a/modules/scripts.py +++ b/modules/scripts.py @@ -500,10 +500,10 @@ class ScriptRunner: s = ScriptSummary('after') script_index = args[0] if len(args) > 0 else 0 if script_index == 0: - return None + return processed script = self.selectable_scripts[script_index-1] if script is None or not hasattr(script, 'after'): - return None + return processed parsed = p.per_script_args.get(script.title(), args[script.args_from:script.args_to]) after_processed = script.after(p, processed, *parsed) if after_processed is not None: diff --git a/scripts/cogvideo.py b/scripts/cogvideo.py new file mode 100644 index 000000000..b8f3ea340 --- /dev/null +++ b/scripts/cogvideo.py @@ -0,0 +1,55 @@ +import os +import gradio as gr +import diffusers +from modules import scripts, processing, shared, devices, sd_models + + +class Script(scripts.Script): + def title(self): + return 'CogVideoX' + + def show(self, is_img2img): + return shared.native + + + def ui(self, _is_img2img): + def video_type_change(video_type): + return [ + gr.update(visible=video_type != 'None'), + gr.update(visible=video_type == 'GIF' or video_type == 'PNG'), + gr.update(visible=video_type == 'MP4'), + gr.update(visible=video_type == 'MP4'), + ] + + 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') + 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) + 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) + 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] + + 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}') + p.extra_generation_params['CogVideoX'] = model + p.do_not_save_grid = True + if 'animatediff' not in p.ops: + p.ops.append('cogvideox') + + 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': + save_video(p, filename=None, images=processed.images, video_type=video_type, duration=duration, loop=loop, pad=pad, interpolate=interpolate)