correct tqdm when inpainting pass is skipped

Author: kabachuha

scripts/deforum/generate.py

@@ -3,6 +3,7 @@
 from PIL import Image, ImageOps
 import requests
 import numpy as np
+from math import ceil
 import torchvision.transforms.functional as TF
 from pytorch_lightning import seed_everything
 import os
@@ -23,6 +24,7 @@
 import cv2
 from .animation import sample_from_cv2, sample_to_cv2
 from modules import processing, masking
+import modules.shared as shared
 from modules.shared import opts, sd_model
 from modules.processing import process_images, StableDiffusionProcessingTxt2Img
 
@@ -185,6 +187,8 @@ def generate(args, root, frame = 0, return_sample=False):
             #color correction for zeroes inpainting
             p.color_corrections = [processing.setup_color_correction(init_image)]
 
+            print("Inpainting zeros")
+
             processed = processing.process_images(p)
 
             init_image = processed.images[0].convert('RGB') 
@@ -193,6 +197,13 @@ def generate(args, root, frame = 0, return_sample=False):
             p.image_mask = None
             mask_image = None
             processed = None
+        else:
+            # fix tqdm total steps if we don't have to conduct a second pass
+            tqdm_instance = shared.total_tqdm
+            current_total = tqdm_instance.getTotal()
+            if current_total != -1:
+                tqdm_instance.updateTotal(current_total - int(ceil(args.steps * (1-args.strength))))
+
 
     elif args.use_init and args.init_image != None and args.init_image != '':
         init_image, mask_image = load_img(args.init_image, 

scripts/deforum/settings.py

@@ -200,3 +200,10 @@ def clear(self):
         if self._tqdm is not None:
             self._tqdm.close()
             self._tqdm = None
+    
+    def getTotal(self):
+        if not opts.multiple_tqdm or cmd_opts.disable_console_progressbars:
+            return -1
+        if self._tqdm is None:
+            self.reset()
+        return self._tqdm.total