File size: 13,319 Bytes
fee99f1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
import os
import platform
import numpy as np
from tqdm import trange
import math
import subprocess as sp
import string
import random
from functools import reduce
import re

import modules.scripts as scripts
import gradio as gr

from modules import processing, shared, sd_samplers, images
from modules.processing import Processed
from modules.sd_samplers import samplers
from modules.shared import opts, cmd_opts, state
import subprocess


wave_completed_regex = r'@wave_completed\(([\-]?[0-9]*\.?[0-9]+), ?([\-]?[0-9]*\.?[0-9]+)\)'
wave_remaining_regex = r'@wave_remaining\(([\-]?[0-9]*\.?[0-9]+), ?([\-]?[0-9]*\.?[0-9]+)\)'

def run_cmd(cmd):
    cmd = list(map(lambda arg: str(arg), cmd))
    print("Executing %s" % " ".join(cmd))
    popen_params = {"stdout": sp.DEVNULL, "stderr": sp.PIPE, "stdin": sp.DEVNULL}

    if os.name == "nt":
       popen_params["creationflags"] = 0x08000000

    proc = sp.Popen(cmd, **popen_params)
    out, err = proc.communicate()  # proc.wait()
    proc.stderr.close()

    if proc.returncode:
        raise IOError(err.decode("utf8"))

    del proc

def encode_video(input_pattern, starting_number, output_dir, fps, quality, encoding, create_segments, segment_duration, ffmpeg_path):
    two_pass = (encoding == "VP9 (webm)")
    alpha_channel = ("webm" in encoding)
    suffix = "webm" if "webm" in encoding else "mp4"
    output_location = output_dir + f".{suffix}"

    encoding_lib = {
      "VP9 (webm)": "libvpx-vp9",
      "VP8 (webm)": "libvpx",
      "H.264 (mp4)": "libx264",
      "H.265 (mp4)": "libx265",
    }[encoding]

    args = [
        "-framerate", fps,
        "-start_number", int(starting_number),
        "-i", input_pattern, 
        "-c:v", encoding_lib, 
        "-b:v","0", 
        "-crf", quality,
        ]

    if encoding_lib == "libvpx-vp9":
        args += ["-pix_fmt", "yuva420p"]
        
    if(ffmpeg_path == ""):
        ffmpeg_path = "ffmpeg"
        if(platform.system == "Windows"):
            ffmpeg_path += ".exe"

    print("\n\n")
    if two_pass:
        first_pass_args = args + [
            "-pass", "1",
            "-an", 
            "-f", "null",
            os.devnull
        ]

        second_pass_args = args + [
            "-pass", "2",
            output_location
        ]

        print("Running first pass ffmpeg encoding")       

        run_cmd([ffmpeg_path] + first_pass_args)
        print("Running second pass ffmpeg encoding.  This could take awhile...")
        run_cmd([ffmpeg_path] + second_pass_args)
    else:
        print("Running ffmpeg encoding.  This could take awhile...")
        run_cmd([ffmpeg_path] + args + [output_location])

    if(create_segments):
      print("Segmenting video")
      run_cmd([ffmpeg_path] + [
          "-i", output_location,
          "-f", "segment",
          "-segment_time", segment_duration,
          "-vcodec", "copy",
          "-acodec", "copy",
          f"{output_dir}.%d.{suffix}"
      ])
      
def set_weights(match_obj, wave_progress):
  weight_0 = 0
  weight_1 = 0
  if match_obj.group(1) is not None:
    weight_0 = float(match_obj.group(1))
  if match_obj.group(2) is not None:
    weight_1 = float(match_obj.group(2))
    
  max_weight = max(weight_0, weight_1)
  min_weight = min(weight_0, weight_1)
  
  weight_range = max_weight - min_weight
  weight = min_weight + weight_range * wave_progress
  return str(weight)


class Script(scripts.Script):
    def title(self):
        return "Loopback Wave V1.4.1"

    def show(self, is_img2img):
        return is_img2img

    def ui(self, is_img2img):
        frames = gr.Slider(minimum=1, maximum=2048, step=1, label='Frames', value=100)
        frames_per_wave = gr.Slider(minimum=0, maximum=120, step=1, label='Frames Per Wave', value=20)
        denoising_strength_change_amplitude = gr.Slider(minimum=0, maximum=1, step=0.01, label='Max additional denoise', value=0.6)
        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)
        initial_image_number = gr.Number(minimum=0, label='Initial generated image number', value=0)

        save_prompts = gr.Checkbox(label='Save prompts as text file', value=True)
        prompts = gr.Textbox(label="Prompt Changes", lines=5, value="")

        save_video = gr.Checkbox(label='Save results as video', value=True)
        output_dir = gr.Textbox(label="Video Name", lines=1, value="")
        video_fps = gr.Slider(minimum=1, maximum=120, step=1, label='Frames per second', value=10)
        video_quality = gr.Slider(minimum=0, maximum=60, step=1, label='Video Quality (crf)', value=40)
        video_encoding = gr.Dropdown(label='Video encoding', value="VP9 (webm)", choices=["VP9 (webm)", "VP8 (webm)", "H.265 (mp4)", "H.264 (mp4)"])
        ffmpeg_path = gr.Textbox(label="ffmpeg binary.  Only set this if it fails otherwise.", lines=1, value="")

        segment_video = gr.Checkbox(label='Cut video in to segments', value=True)
        video_segment_duration = gr.Slider(minimum=10, maximum=60, step=1, label='Video Segment Duration (seconds)', value=20)


        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]

    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):
        processing.fix_seed(p)
        batch_count = p.n_iter
        p.extra_generation_params = {
            "Max Additional Denoise": denoising_strength_change_amplitude,
            "Frames per wave": frames_per_wave,
            "Wave Offset": denoising_strength_change_offset,
        }

        # We save them ourselves for the sake of ffmpeg
        p.do_not_save_samples = True

        changes_dict = {}


        p.batch_size = 1
        p.n_iter = 1

        output_images, info = None, None
        initial_seed = None
        initial_info = None

        grids = []
        all_images = []
        original_init_image = p.init_images
        state.job_count = frames * batch_count

        initial_color_corrections = [processing.setup_color_correction(p.init_images[0])]
        initial_denoising_strength = p.denoising_strength

        if(output_dir==""):
            output_dir = str(p.seed)
        else:
            output_dir = output_dir + "-" + str(p.seed)

        loopback_wave_path = os.path.join(p.outpath_samples, "loopback-wave")
        loopback_wave_images_path = os.path.join(loopback_wave_path, output_dir)

        os.makedirs(loopback_wave_images_path, exist_ok=True)

        p.outpath_samples = loopback_wave_images_path
        
        prompts = prompts.strip()
        
        if save_prompts:
            with open(loopback_wave_images_path + "-prompts.txt", "w") as f:
                generation_settings = [
                  "Generation Settings",
                  f"Total Frames: {frames}",
                  f"Frames Per Wave: {frames_per_wave}",
                  f"Wave Offset: {denoising_strength_change_offset}",
                  f"Base Denoising Strength: {initial_denoising_strength}",
                  f"Max Additional Denoise: {denoising_strength_change_amplitude}",
                  f"Initial Image Number: {initial_image_number}",
                  "",
                  "Video Encoding Settings",
                  f"Save Video: {save_video}"
                ]
                
                if save_video:
                  generation_settings = generation_settings + [
                    f"Framerate: {video_fps}",
                    f"Quality: {video_quality}",
                    f"Encoding: {video_encoding}",
                    f"Create Segmented Video: {segment_video}"
                  ]
                  
                  if segment_video:
                    generation_settings = generation_settings + [f"Segment Duration: {video_segment_duration}"]
                 
                generation_settings = generation_settings + [
                  "",
                  "Prompt Details",
                  "Initial Prompt:" + p.prompt,
                  "",
                  "Negative Prompt:" + p.negative_prompt,
                  "",
                  "Frame change prompts:",
                  prompts
                ]
                  


                f.write('\n'.join(generation_settings))

        if prompts:
            lines = prompts.split("\n")
            for prompt_line in lines:
              params = prompt_line.split("::")
              if len(params) == 2:
                changes_dict[params[0]] = { "prompt": params[1] }
              elif len(params) == 3:
                changes_dict[params[0]] = { "seed": params[1], "prompt": params[2] }
              else:
                raise IOError(f"Invalid input in prompt line: {prompt_line}")
        
        raw_prompt = p.prompt
                
        for n in range(batch_count):
            history = []

            # Reset to original init image at the start of each batch
            p.init_images = original_init_image
            
            seed_state = "adding"
            current_seed = p.seed

            for i in range(frames):
                current_seed = p.seed
                state.job = ""
                
                if str(i) in changes_dict:
                  raw_prompt = changes_dict[str(i)]["prompt"]
                  state.job = "New prompt: %s\n" % raw_prompt
                                    
                  if "seed" in changes_dict[str(i)]:
                    current_seed = changes_dict[str(i)]["seed"]
                    
                    if current_seed.startswith("+"):
                      seed_state = "adding"
                      current_seed = current_seed.strip("+")
                    elif current_seed.startswith("-"):
                      seed_state = "subtracting"
                      current_seed = current_seed.strip("-")
                    else:
                      seed_state = "constant"
                      
                    current_seed = int(current_seed)
                    p.seed = current_seed
                      
                      
                  
                p.n_iter = 1
                p.batch_size = 1
                p.do_not_save_grid = True

                if opts.img2img_color_correction:
                    p.color_corrections = initial_color_corrections
                    
                    
                wave_progress = float(1)/(float(frames_per_wave - 1))*float(((float(i)%float(frames_per_wave)) + ((float(1)/float(180))*denoising_strength_change_offset)))
                print(wave_progress)
                new_prompt = re.sub(wave_completed_regex, lambda x: set_weights(x, wave_progress), raw_prompt)
                new_prompt = re.sub(wave_remaining_regex, lambda x: set_weights(x, 1 - wave_progress), new_prompt)
                p.prompt = new_prompt
                
                print(new_prompt)

                denoising_strength_change_rate = 180/frames_per_wave

                cos = abs(math.cos(math.radians(i*denoising_strength_change_rate + denoising_strength_change_offset)))
                p.denoising_strength = initial_denoising_strength + denoising_strength_change_amplitude - (cos * denoising_strength_change_amplitude)

                state.job += f"Iteration {i + 1}/{frames}, batch {n + 1}/{batch_count}. Denoising Strength: {p.denoising_strength}"

                processed = processing.process_images(p)

                if initial_seed is None:
                    initial_seed = processed.seed
                    initial_info = processed.info

                init_img = processed.images[0]

                p.init_images = [init_img]
                
                if seed_state == "adding":
                  p.seed = processed.seed + 1
                elif seed_state == "subtracting":
                  p.seed = processed.seed - 1
                  
                image_number = int(initial_image_number + i)
                images.save_image(init_img, p.outpath_samples, "", processed.seed, processed.prompt, forced_filename=str(image_number))

                history.append(init_img)

            grid = images.image_grid(history, rows=1)
            if opts.grid_save:
                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)

            grids.append(grid)
            all_images += history

        if opts.return_grid:
            all_images = grids + all_images

        if save_video:
            input_pattern = os.path.join(loopback_wave_images_path, "%d.png")
            encode_video(input_pattern, initial_image_number, loopback_wave_images_path, video_fps, video_quality, video_encoding, segment_video, video_segment_duration, ffmpeg_path)

        processed = Processed(p, all_images, initial_seed, initial_info)

        return processed