Upload 3 files

app.py

@@ -167,7 +167,7 @@ from modutils import (list_uniq, download_private_repo, get_model_id_list, get_t
     get_my_lora, upload_file_lora, move_file_lora, search_civitai_lora, select_civitai_lora,
     set_textual_inversion_prompt, get_model_pipeline, change_interface_mode, get_t2i_model_info,
     get_tupled_model_list, save_gallery_images, set_optimization, set_sampler_settings,
-    set_quick_presets, process_style_prompt, optimization_list,
+    set_quick_presets, process_style_prompt, optimization_list, save_images,
     preset_styles, preset_quality, preset_sampler_setting, translate_to_en)
 from env import (hf_token, CIVITAI_API_KEY, HF_LORA_ESSENTIAL_PRIVATE_REPO, HF_VAE_PRIVATE_REPO,
     HF_SDXL_EMBEDS_NEGATIVE_PRIVATE_REPO, HF_SDXL_EMBEDS_POSITIVE_PRIVATE_REPO,
@@ -556,16 +556,17 @@ class GuiSD:
         model_precision = torch.float16
         self.model.device = torch.device("cuda:0")
         if not self.model:
-            from stablepy import Model_Diffusers
+            from modelstream import Model_Diffusers2
 
             print("Loading model...")
-            self.model = Model_Diffusers(
+            self.model = Model_Diffusers2(
                 base_model_id=model_name,
                 task_name=task,
                 vae_model=vae_model if vae_model != "None" else None,
                 type_model_precision=model_precision,
                 retain_task_model_in_cache=retain_task_cache_gui,
             )
+        self.model.stream_config(concurrency=5, latent_resize_by=1, vae_decoding=False)
 
         if task != "txt2img" and not image_control:
             raise ValueError("No control image found: To use this function, you have to upload an image in 'Image ControlNet/Inpaint/Img2img'")
@@ -710,18 +711,24 @@ class GuiSD:
         }
 
         # Maybe fix lora issue: 'Cannot copy out of meta tensor; no data!''
-        self.model.pipe.to("cuda:0" if torch.cuda.is_available() else "cpu")
+        #self.model.pipe.to("cuda:0" if torch.cuda.is_available() else "cpu")
 
         #progress(0, desc="Preparation completed. Starting inference...")
+
         info_state = f"PROCESSING "
-        for img, seed, data in self.model(**pipe_params):
+        for img, seed, image_path, metadata in self.model(**pipe_params):
             info_state += ">"
-            if data:
+            if image_path:
                 info_state = f"COMPLETED. Seeds: {str(seed)}"
                 if vae_msg:
                     info_state = info_state + "<br>" + vae_msg
                 if msg_lora:
                     info_state = info_state + "<br>" + "<br>".join(msg_lora)
+
+                info_state = info_state + "<br>" + "GENERATION DATA:<br>" + "<br>-------<br>".join(metadata).replace("\n", "<br>")
+
+                img = save_images(img, metadata)
+                
             yield img, info_state
 
 sd_gen = GuiSD()

modutils.py

@@ -100,6 +100,24 @@ def safe_float(input):
     return output
 
 
+from PIL import Image
+def save_images(images: list[Image.Image], metadatas: list[str]):
+    from PIL import PngImagePlugin
+    import uuid
+    try:
+        output_images = []
+        for image, metadata in zip(images, metadatas):
+            info = PngImagePlugin.PngInfo()
+            info.add_text("metadata", metadata)
+            savefile = f"{str(uuid.uuid4())}.png"
+            image.save(savefile, "PNG", pnginfo=info)
+            output_images.append(str(Path(savefile).resolve()))
+        return output_images
+    except Exception as e:
+        print(f"Failed to save image file: {e}")
+        raise Exception(f"Failed to save image file:") from e
+
+
 def save_gallery_images(images, progress=gr.Progress(track_tqdm=True)):
     from datetime import datetime, timezone, timedelta
     progress(0, desc="Updating gallery...")

requirements.txt

@@ -2,6 +2,7 @@ git+https://github.com/R3gm/stablepy.git
 torch==2.2.0
 gdown
 opencv-python
+yt-dlp
 torchvision
 accelerate
 transformers