hvai/sdset

1import os2import platform3import numpy as np4from tqdm import trange5import math6import subprocess as sp7import string8import random9from functools import reduce10import re11 12import modules.scripts as scripts13import gradio as gr14 15from modules import processing, shared, sd_samplers, images16from modules.processing import Processed17from modules.sd_samplers import samplers18from modules.shared import opts, cmd_opts, state19import subprocess20 21 22wave_completed_regex = r'@wave_completed\(([\-]?[0-9]*\.?[0-9]+), ?([\-]?[0-9]*\.?[0-9]+)\)'23wave_remaining_regex = r'@wave_remaining\(([\-]?[0-9]*\.?[0-9]+), ?([\-]?[0-9]*\.?[0-9]+)\)'24 25def run_cmd(cmd):26    cmd = list(map(lambda arg: str(arg), cmd))27    print("Executing %s" % " ".join(cmd))28    popen_params = {"stdout": sp.DEVNULL, "stderr": sp.PIPE, "stdin": sp.DEVNULL}29 30    if os.name == "nt":31       popen_params["creationflags"] = 0x0800000032 33    proc = sp.Popen(cmd, **popen_params)34    out, err = proc.communicate()  # proc.wait()35    proc.stderr.close()36 37    if proc.returncode:38        raise IOError(err.decode("utf8"))39 40    del proc41 42def encode_video(input_pattern, starting_number, output_dir, fps, quality, encoding, create_segments, segment_duration, ffmpeg_path):43    two_pass = (encoding == "VP9 (webm)")44    alpha_channel = ("webm" in encoding)45    suffix = "webm" if "webm" in encoding else "mp4"46    output_location = output_dir + f".{suffix}"47 48    encoding_lib = {49      "VP9 (webm)": "libvpx-vp9",50      "VP8 (webm)": "libvpx",51      "H.264 (mp4)": "libx264",52      "H.265 (mp4)": "libx265",53    }[encoding]54 55    args = [56        "-framerate", fps,57        "-start_number", int(starting_number),58        "-i", input_pattern, 59        "-c:v", encoding_lib, 60        "-b:v","0", 61        "-crf", quality,62        ]63 64    if encoding_lib == "libvpx-vp9":65        args += ["-pix_fmt", "yuva420p"]66        67    if(ffmpeg_path == ""):68        ffmpeg_path = "ffmpeg"69        if(platform.system == "Windows"):70            ffmpeg_path += ".exe"71 72    print("\n\n")73    if two_pass:74        first_pass_args = args + [75            "-pass", "1",76            "-an", 77            "-f", "null",78            os.devnull79        ]80 81        second_pass_args = args + [82            "-pass", "2",83            output_location84        ]85 86        print("Running first pass ffmpeg encoding")       87 88        run_cmd([ffmpeg_path] + first_pass_args)89        print("Running second pass ffmpeg encoding.  This could take awhile...")90        run_cmd([ffmpeg_path] + second_pass_args)91    else:92        print("Running ffmpeg encoding.  This could take awhile...")93        run_cmd([ffmpeg_path] + args + [output_location])94 95    if(create_segments):96      print("Segmenting video")97      run_cmd([ffmpeg_path] + [98          "-i", output_location,99          "-f", "segment",100          "-segment_time", segment_duration,101          "-vcodec", "copy",102          "-acodec", "copy",103          f"{output_dir}.%d.{suffix}"104      ])105      106def set_weights(match_obj, wave_progress):107  weight_0 = 0108  weight_1 = 0109  if match_obj.group(1) is not None:110    weight_0 = float(match_obj.group(1))111  if match_obj.group(2) is not None:112    weight_1 = float(match_obj.group(2))113    114  max_weight = max(weight_0, weight_1)115  min_weight = min(weight_0, weight_1)116  117  weight_range = max_weight - min_weight118  weight = min_weight + weight_range * wave_progress119  return str(weight)120 121 122class Script(scripts.Script):123    def title(self):124        return "Loopback Wave V1.4.1"125 126    def show(self, is_img2img):127        return is_img2img128 129    def ui(self, is_img2img):130        frames = gr.Slider(minimum=1, maximum=2048, step=1, label='Frames', value=100)131        frames_per_wave = gr.Slider(minimum=0, maximum=120, step=1, label='Frames Per Wave', value=20)132        denoising_strength_change_amplitude = gr.Slider(minimum=0, maximum=1, step=0.01, label='Max additional denoise', value=0.6)133        denoising_strength_change_offset = gr.Number(minimum=0, maximum=180, step=1, label='Wave offset (ignore this if you don\'t know what it means)', value=0)134        initial_image_number = gr.Number(minimum=0, label='Initial generated image number', value=0)135 136        save_prompts = gr.Checkbox(label='Save prompts as text file', value=True)137        prompts = gr.Textbox(label="Prompt Changes", lines=5, value="")138 139        save_video = gr.Checkbox(label='Save results as video', value=True)140        output_dir = gr.Textbox(label="Video Name", lines=1, value="")141        video_fps = gr.Slider(minimum=1, maximum=120, step=1, label='Frames per second', value=10)142        video_quality = gr.Slider(minimum=0, maximum=60, step=1, label='Video Quality (crf)', value=40)143        video_encoding = gr.Dropdown(label='Video encoding', value="VP9 (webm)", choices=["VP9 (webm)", "VP8 (webm)", "H.265 (mp4)", "H.264 (mp4)"])144        ffmpeg_path = gr.Textbox(label="ffmpeg binary.  Only set this if it fails otherwise.", lines=1, value="")145 146        segment_video = gr.Checkbox(label='Cut video in to segments', value=True)147        video_segment_duration = gr.Slider(minimum=10, maximum=60, step=1, label='Video Segment Duration (seconds)', value=20)148 149 150        return [frames, denoising_strength_change_amplitude, frames_per_wave, denoising_strength_change_offset,initial_image_number, prompts, save_prompts, save_video, output_dir, video_fps, video_quality, video_encoding, ffmpeg_path, segment_video, video_segment_duration]151 152    def run(self, p, frames, denoising_strength_change_amplitude, frames_per_wave, denoising_strength_change_offset, initial_image_number, prompts: str,save_prompts, save_video, output_dir, video_fps, video_quality, video_encoding, ffmpeg_path, segment_video, video_segment_duration):153        processing.fix_seed(p)154        batch_count = p.n_iter155        p.extra_generation_params = {156            "Max Additional Denoise": denoising_strength_change_amplitude,157            "Frames per wave": frames_per_wave,158            "Wave Offset": denoising_strength_change_offset,159        }160 161        # We save them ourselves for the sake of ffmpeg162        p.do_not_save_samples = True163 164        changes_dict = {}165 166 167        p.batch_size = 1168        p.n_iter = 1169 170        output_images, info = None, None171        initial_seed = None172        initial_info = None173 174        grids = []175        all_images = []176        original_init_image = p.init_images177        state.job_count = frames * batch_count178 179        initial_color_corrections = [processing.setup_color_correction(p.init_images[0])]180        initial_denoising_strength = p.denoising_strength181 182        if(output_dir==""):183            output_dir = str(p.seed)184        else:185            output_dir = output_dir + "-" + str(p.seed)186 187        loopback_wave_path = os.path.join(p.outpath_samples, "loopback-wave")188        loopback_wave_images_path = os.path.join(loopback_wave_path, output_dir)189 190        os.makedirs(loopback_wave_images_path, exist_ok=True)191 192        p.outpath_samples = loopback_wave_images_path193        194        prompts = prompts.strip()195        196        if save_prompts:197            with open(loopback_wave_images_path + "-prompts.txt", "w") as f:198                generation_settings = [199                  "Generation Settings",200                  f"Total Frames: {frames}",201                  f"Frames Per Wave: {frames_per_wave}",202                  f"Wave Offset: {denoising_strength_change_offset}",203                  f"Base Denoising Strength: {initial_denoising_strength}",204                  f"Max Additional Denoise: {denoising_strength_change_amplitude}",205                  f"Initial Image Number: {initial_image_number}",206                  "",207                  "Video Encoding Settings",208                  f"Save Video: {save_video}"209                ]210                211                if save_video:212                  generation_settings = generation_settings + [213                    f"Framerate: {video_fps}",214                    f"Quality: {video_quality}",215                    f"Encoding: {video_encoding}",216                    f"Create Segmented Video: {segment_video}"217                  ]218                  219                  if segment_video:220                    generation_settings = generation_settings + [f"Segment Duration: {video_segment_duration}"]221                 222                generation_settings = generation_settings + [223                  "",224                  "Prompt Details",225                  "Initial Prompt:" + p.prompt,226                  "",227                  "Negative Prompt:" + p.negative_prompt,228                  "",229                  "Frame change prompts:",230                  prompts231                ]232                  233 234 235                f.write('\n'.join(generation_settings))236 237        if prompts:238            lines = prompts.split("\n")239            for prompt_line in lines:240              params = prompt_line.split("::")241              if len(params) == 2:242                changes_dict[params[0]] = { "prompt": params[1] }243              elif len(params) == 3:244                changes_dict[params[0]] = { "seed": params[1], "prompt": params[2] }245              else:246                raise IOError(f"Invalid input in prompt line: {prompt_line}")247        248        raw_prompt = p.prompt249                250        for n in range(batch_count):251            history = []252 253            # Reset to original init image at the start of each batch254            p.init_images = original_init_image255            256            seed_state = "adding"257            current_seed = p.seed258 259            for i in range(frames):260                current_seed = p.seed261                state.job = ""262                263                if str(i) in changes_dict:264                  raw_prompt = changes_dict[str(i)]["prompt"]265                  state.job = "New prompt: %s\n" % raw_prompt266                                    267                  if "seed" in changes_dict[str(i)]:268                    current_seed = changes_dict[str(i)]["seed"]269                    270                    if current_seed.startswith("+"):271                      seed_state = "adding"272                      current_seed = current_seed.strip("+")273                    elif current_seed.startswith("-"):274                      seed_state = "subtracting"275                      current_seed = current_seed.strip("-")276                    else:277                      seed_state = "constant"278                      279                    current_seed = int(current_seed)280                    p.seed = current_seed281                      282                      283                  284                p.n_iter = 1285                p.batch_size = 1286                p.do_not_save_grid = True287 288                if opts.img2img_color_correction:289                    p.color_corrections = initial_color_corrections290                    291                    292                wave_progress = float(1)/(float(frames_per_wave - 1))*float(((float(i)%float(frames_per_wave)) + ((float(1)/float(180))*denoising_strength_change_offset)))293                print(wave_progress)294                new_prompt = re.sub(wave_completed_regex, lambda x: set_weights(x, wave_progress), raw_prompt)295                new_prompt = re.sub(wave_remaining_regex, lambda x: set_weights(x, 1 - wave_progress), new_prompt)296                p.prompt = new_prompt297                298                print(new_prompt)299 300                denoising_strength_change_rate = 180/frames_per_wave301 302                cos = abs(math.cos(math.radians(i*denoising_strength_change_rate + denoising_strength_change_offset)))303                p.denoising_strength = initial_denoising_strength + denoising_strength_change_amplitude - (cos * denoising_strength_change_amplitude)304 305                state.job += f"Iteration {i + 1}/{frames}, batch {n + 1}/{batch_count}. Denoising Strength: {p.denoising_strength}"306 307                processed = processing.process_images(p)308 309                if initial_seed is None:310                    initial_seed = processed.seed311                    initial_info = processed.info312 313                init_img = processed.images[0]314 315                p.init_images = [init_img]316                317                if seed_state == "adding":318                  p.seed = processed.seed + 1319                elif seed_state == "subtracting":320                  p.seed = processed.seed - 1321                  322                image_number = int(initial_image_number + i)323                images.save_image(init_img, p.outpath_samples, "", processed.seed, processed.prompt, forced_filename=str(image_number))324 325                history.append(init_img)326 327            grid = images.image_grid(history, rows=1)328            if opts.grid_save:329                images.save_image(grid, p.outpath_grids, "grid", initial_seed, p.prompt, opts.grid_format, info=info, short_filename=not opts.grid_extended_filename, grid=True, p=p)330 331            grids.append(grid)332            all_images += history333 334        if opts.return_grid:335            all_images = grids + all_images336 337        if save_video:338            input_pattern = os.path.join(loopback_wave_images_path, "%d.png")339            encode_video(input_pattern, initial_image_number, loopback_wave_images_path, video_fps, video_quality, video_encoding, segment_video, video_segment_duration, ffmpeg_path)340 341        processed = Processed(p, all_images, initial_seed, initial_info)342 343        return processed344 345 346