Restore zeros

app.py

@@ -455,8 +455,6 @@ def worker(input_image, image_position, prompts, n_prompt, seed, resolution, tot
             clean_latent_indices = torch.cat([clean_latent_1x_indices, clean_latent_indices_start], dim=1)
 
         def post_process(forward, generated_latents, total_generated_latent_frames, history_latents, high_vram, transformer, gpu, vae, history_pixels, latent_window_size, enable_preview, section_index, total_latent_sections, outputs_folder, mp4_crf, stream):
-            if int(history_latents.shape[2]) > total_generated_latent_frames:
-                history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :] if forward else history_latents[:, :, :total_generated_latent_frames, :, :]
             total_generated_latent_frames += int(generated_latents.shape[2])
             history_latents = torch.cat([history_latents, generated_latents.to(history_latents)] if forward else [generated_latents.to(history_latents), history_latents], dim=2)
 
@@ -465,15 +463,18 @@ def worker(input_image, image_position, prompts, n_prompt, seed, resolution, tot
                 load_model_as_complete(vae, target_device=gpu)
 
             if history_pixels is None:
-                history_pixels = vae_decode(history_latents, vae).cpu()
+                real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :] if forward else history_latents[:, :, :total_generated_latent_frames, :, :]
+                history_pixels = vae_decode(real_history_latents, vae).cpu()
             else:
                 section_latent_frames = latent_window_size * 2
                 overlapped_frames = latent_window_size * 4 - 3
 
                 if forward:
-                    history_pixels = soft_append_bcthw(history_pixels, vae_decode(history_latents[:, :, -section_latent_frames:, :, :], vae).cpu(), overlapped_frames)
+                    real_history_latents = history_latents[:, :, -min(section_latent_frames, total_generated_latent_frames):, :, :]
+                    history_pixels = soft_append_bcthw(history_pixels, vae_decode(real_history_latents, vae).cpu(), overlapped_frames)
                 else:
-                    history_pixels = soft_append_bcthw(vae_decode(history_latents[:, :, :section_latent_frames, :, :], vae).cpu(), history_pixels, overlapped_frames)
+                    real_history_latents = history_latents[:, :, :min(section_latent_frames, total_generated_latent_frames), :, :]
+                    history_pixels = soft_append_bcthw(vae_decode(real_history_latents, vae).cpu(), history_pixels, overlapped_frames)
 
             if not high_vram:
                 unload_complete_models(text_encoder, text_encoder_2, image_encoder, vae, transformer)
@@ -557,6 +558,11 @@ def worker(input_image, image_position, prompts, n_prompt, seed, resolution, tot
                 else:
                     clean_latent_indices_start, clean_latent_4x_indices, clean_latent_2x_indices, clean_latent_1x_indices, latent_indices = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
                     clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
+
+                    real_history_latents = history_latents[:, :, :total_generated_latent_frames, :, :]
+                    zero_latents = history_latents[:, :, total_generated_latent_frames:, :, :]
+                    history_latents = torch.cat([zero_latents, real_history_latents], dim=2)
+                    real_history_latents = zero_latents = None
                     
                     forward = True
                     section_index = first_section_index
@@ -1084,7 +1090,7 @@ title_html = """
 js = """
 function createGradioAnimation() {
     window.addEventListener("beforeunload", function(e) {
-        if (document.getElementById('warning') && !document.getElementById('warning').disabled) {
+        if (document.getElementById('end-button') && !document.getElementById('end-button').disabled) {
             var confirmationMessage = 'A process is still running. '
                                     + 'If you leave before saving, your changes will be lost.';