Upload dc.py

dc.py

@@ -501,9 +501,51 @@ class GuiSD:
             self.model.pipe.transformer.to(self.model.device)
             print("transformer to cuda")
 
-        return self.infer_short(self.model, pipe_params), info_state
+        #return self.infer_short(self.model, pipe_params), info_state
+
+        actual_progress = 0
+        info_images = gr.update()
+        for img, seed, image_path, metadata in self.model(**pipe_params):
+            info_state = progress_step_bar(actual_progress, steps)
+            actual_progress += concurrency
+            if image_path:
+                info_images = f"Seeds: {str(seed)}"
+                if vae_msg:
+                    info_images = info_images + "<br>" + vae_msg
+
+                if "Cannot copy out of meta tensor; no data!" in self.model.last_lora_error:
+                    msg_ram = "Unable to process the LoRAs due to high RAM usage; please try again later."
+                    print(msg_ram)
+                    msg_lora += f"<br>{msg_ram}"
+
+                for status, lora in zip(self.model.lora_status, self.model.lora_memory):
+                    if status:
+                        msg_lora += f"<br>Loaded: {lora}"
+                    elif status is not None:
+                        msg_lora += f"<br>Error with: {lora}"
+
+                if msg_lora:
+                    info_images += msg_lora
+
+                info_images = info_images + "<br>" + "GENERATION DATA:<br>" + metadata[0].replace("\n", "<br>") + "<br>-------<br>"
+
+                download_links = "<br>".join(
+                    [
+                        f'<a href="{path.replace("/images/", "/file=/home/user/app/images/")}" download="{os.path.basename(path)}">Download Image {i + 1}</a>'
+                        for i, path in enumerate(image_path)
+                    ]
+                )
+                if save_generated_images:
+                    info_images += f"<br>{download_links}"
+
+## BEGIN MOD
+                img = save_images(img, metadata)
 ## END MOD
 
+                info_state = "COMPLETE"
+
+            #yield info_state, img, info_images
+            return img, info_state
 
 def dynamic_gpu_duration(func, duration, *args):