Progress bar for loading pipeline

app.py

@@ -14,9 +14,9 @@ filterwarnings("ignore", category=FutureWarning, module="diffusers")
 filterwarnings("ignore", category=FutureWarning, module="transformers")
 
 diffusers_logging.set_verbosity_error()
-diffusers_logging.disable_progress_bar()
-
 transformers_logging.set_verbosity_error()
+
+diffusers_logging.disable_progress_bar()
 transformers_logging.disable_progress_bar()
 
 # the CSS `content` attribute expects a string so we need to wrap the number in quotes
@@ -88,7 +88,7 @@ async def random_fn():
     return gr.Textbox(value=random.choice(prompts))
 
 
-async def generate_fn(*args):
+async def generate_fn(*args, progress=gr.Progress(track_tqdm=True)):
     if len(args) > 0:
         prompt = args[0]
     else:
@@ -104,12 +104,15 @@ async def generate_fn(*args):
         gen_args[3] = None
 
     try:
+        if Config.ZERO_GPU:
+            progress((0, 100), desc="ZeroGPU init")
+
         images = await async_call(
             generate,
             *gen_args,
-            Info=gr.Info,
             Error=gr.Error,
-            Progress=gr.Progress,
+            Info=gr.Info,
+            progress=progress,
         )
     except RuntimeError:
         raise gr.Error("Error: Please try again")

cli.py

@@ -2,9 +2,20 @@
 # usage: python cli.py 'colorful calico cat artstation'
 import argparse
 import asyncio
+from warnings import filterwarnings
+
+from diffusers.utils import logging as diffusers_logging
+from transformers import logging as transformers_logging
 
 from lib import Config, async_call, generate
 
+filterwarnings("ignore", category=FutureWarning, module="diffusers")
+filterwarnings("ignore", category=FutureWarning, module="transformers")
+
+# reduce verbosity but don't disable progress bars
+diffusers_logging.set_verbosity_error()
+transformers_logging.set_verbosity_error()
+
 
 def save_images(images, filename="image.png"):
     for i, (img, _) in enumerate(images):

lib/config.py

@@ -1,4 +1,5 @@
 import os
+from importlib import import_module
 from types import SimpleNamespace
 
 from diffusers import (
@@ -8,14 +9,18 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
     PNDMScheduler,
-    StableDiffusionImg2ImgPipeline,
-    StableDiffusionPipeline,
     UniPCMultistepScheduler,
 )
 
+from .pipelines import CustomStableDiffusionImg2ImgPipeline, CustomStableDiffusionPipeline
+
+# improved GPU handling and progress bars; set before importing spaces
+os.environ["ZEROGPU_V2"] = "true"
+
 Config = SimpleNamespace(
     HF_TOKEN=os.environ.get("HF_TOKEN", None),
     CIVIT_TOKEN=os.environ.get("CIVIT_TOKEN", None),
+    ZERO_GPU=import_module("spaces").config.Config.zero_gpu,
     HF_MODELS={
         "Lykon/dreamshaper-8": [
             "feature_extractor/preprocessor_config.json",
@@ -59,8 +64,8 @@ Config = SimpleNamespace(
         "Noto Color Emoji",
     ],
     PIPELINES={
-        "txt2img": StableDiffusionPipeline,
-        "img2img": StableDiffusionImg2ImgPipeline,
+        "txt2img": CustomStableDiffusionPipeline,
+        "img2img": CustomStableDiffusionImg2ImgPipeline,
     },
     MODEL="Lykon/dreamshaper-8",
     MODELS=[

lib/inference.py

@@ -4,14 +4,13 @@ import time
 from datetime import datetime
 from itertools import product
 
-import gradio as gr
 import numpy as np
-import spaces
 import torch
 from compel import Compel, DiffusersTextualInversionManager, ReturnedEmbeddingsType
 from compel.prompt_parser import PromptParser
 from huggingface_hub.utils import HFValidationError, RepositoryNotFoundError
 from PIL import Image
+from spaces import GPU
 
 from .config import Config
 from .loader import Loader
@@ -92,7 +91,7 @@ def gpu_duration(**kwargs):
     return loading + (duration * num_images)
 
 
-@spaces.GPU(duration=gpu_duration)
+@GPU(duration=gpu_duration)
 def generate(
     positive_prompt,
     negative_prompt="",
@@ -120,10 +119,9 @@ def generate(
     taesd=False,
     freeu=False,
     clip_skip=False,
-    Info=None,
     Error=Exception,
-    Progress=None,
-    progress=gr.Progress(track_tqdm=True),
+    Info=None,
+    progress=None,
 ):
     if not torch.cuda.is_available():
         raise Error("CUDA not available")
@@ -148,32 +146,27 @@ def generate(
     else:
         IP_ADAPTER = ""
 
-    if Progress is not None:
-        TQDM = False
-        progress_bar = Progress()
-        progress_bar((0, inference_steps), desc=f"Generating image {CURRENT_IMAGE}/{num_images}")
-    else:
-        TQDM = True
-        progress_bar = None
-
+    # custom progress bar for multiple images
     def callback_on_step_end(pipeline, step, timestep, latents):
         nonlocal CURRENT_STEP, CURRENT_IMAGE
-        if Progress is None:
-            return latents
-        strength = denoising_strength if KIND == "img2img" else 1
-        total_steps = min(int(inference_steps * strength), inference_steps)
-
-        CURRENT_STEP = step + 1
-        progress_bar(
-            (CURRENT_STEP, total_steps),
-            desc=f"Generating image {CURRENT_IMAGE}/{num_images}",
-        )
+        if progress is not None:
+            # calculate total steps for img2img based on denoising strength
+            strength = denoising_strength if KIND == "img2img" else 1
+            total_steps = min(int(inference_steps * strength), inference_steps)
+            CURRENT_STEP = step + 1
+            progress(
+                (CURRENT_STEP, total_steps),
+                desc=f"Generating image {CURRENT_IMAGE}/{num_images}",
+            )
         return latents
 
     start = time.perf_counter()
     log = Logger("generate")
     log.info(f"Generating {num_images} image{'s' if num_images > 1 else ''}")
 
+    if Config.ZERO_GPU and progress is not None:
+        progress((100, 100), desc="ZeroGPU init")
+
     loader = Loader()
     loader.load(
         KIND,
@@ -185,7 +178,7 @@ def generate(
         freeu,
         deepcache,
         scale,
-        TQDM,
+        progress,
     )
 
     if loader.pipe is None:

lib/loader.py

@@ -4,7 +4,6 @@ from threading import Lock
 
 import torch
 from DeepCache import DeepCacheSDHelper
-from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionPipeline
 from diffusers.models import AutoencoderKL, AutoencoderTiny
 from diffusers.models.attention_processor import AttnProcessor2_0, IPAdapterAttnProcessor2_0
 
@@ -50,9 +49,9 @@ class Loader:
             return False
         if self.model.lower() != model.lower():
             return True
-        if kind == "txt2img" and not isinstance(self.pipe, StableDiffusionPipeline):
+        if kind == "txt2img" and not isinstance(self.pipe, Config.PIPELINES["txt2img"]):
             return True  # txt2img -> img2img
-        if kind == "img2img" and not isinstance(self.pipe, StableDiffusionImg2ImgPipeline):
+        if kind == "img2img" and not isinstance(self.pipe, Config.PIPELINES["img2img"]):
             return True  # img2img -> txt2img
         return False
 
@@ -69,7 +68,7 @@ class Loader:
             return
 
         self.log.info("Unloading IP-Adapter")
-        if not isinstance(self.pipe, StableDiffusionImg2ImgPipeline):
+        if not isinstance(self.pipe, Config.PIPELINES["img2img"]):
             self.pipe.image_encoder = None
             self.pipe.register_to_config(image_encoder=[None, None])
 
@@ -142,7 +141,13 @@ class Loader:
                 self.log.error(f"Error loading 4x upscaler: {e}")
                 self.upscaler_4x = None
 
-    def _load_pipeline(self, kind, model, tqdm, **kwargs):
+    def _load_pipeline(
+        self,
+        kind,
+        model,
+        progress,
+        **kwargs,
+    ):
         pipeline = Config.PIPELINES[kind]
         if self.pipe is None:
             try:
@@ -152,10 +157,11 @@ class Loader:
                 if model.lower() in Config.MODEL_CHECKPOINTS.keys():
                     self.pipe = pipeline.from_single_file(
                         f"https://huggingface.co/{model}/{Config.MODEL_CHECKPOINTS[model.lower()]}",
+                        progress,
                         **kwargs,
                     ).to("cuda")
                 else:
-                    self.pipe = pipeline.from_pretrained(model, **kwargs).to("cuda")
+                    self.pipe = pipeline.from_pretrained(model, progress, **kwargs).to("cuda")
                 diff = time.perf_counter() - start
                 self.log.info(f"Loading {model} done in {diff:.2f}s")
             except Exception as e:
@@ -166,7 +172,7 @@ class Loader:
         if not isinstance(self.pipe, pipeline):
             self.pipe = pipeline.from_pipe(self.pipe).to("cuda")
         if self.pipe is not None:
-            self.pipe.set_progress_bar_config(disable=not tqdm)
+            self.pipe.set_progress_bar_config(disable=progress is not None)
 
     def _load_vae(self, taesd=False, model=""):
         vae_type = type(self.pipe.vae)
@@ -231,7 +237,7 @@ class Loader:
         freeu,
         deepcache,
         scale,
-        tqdm,
+        progress,
     ):
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -275,7 +281,7 @@ class Loader:
             pipe_kwargs["torch_dtype"] = torch.float16
 
         self._unload(kind, model, ip_adapter, deepcache)
-        self._load_pipeline(kind, model, tqdm, **pipe_kwargs)
+        self._load_pipeline(kind, model, progress, **pipe_kwargs)
 
         # error loading model
         if self.pipe is None:

lib/pipelines.py

@@ -0,0 +1,222 @@
+import os
+from importlib import import_module
+
+from diffusers import StableDiffusionImg2ImgPipeline, StableDiffusionPipeline
+from diffusers.loaders.single_file import (
+    SINGLE_FILE_OPTIONAL_COMPONENTS,
+    load_single_file_sub_model,
+)
+from diffusers.loaders.single_file_utils import fetch_diffusers_config, load_single_file_checkpoint
+from diffusers.models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT
+from diffusers.pipelines.pipeline_loading_utils import (
+    ALL_IMPORTABLE_CLASSES,
+    _get_pipeline_class,
+    load_sub_model,
+)
+from diffusers.utils import logging
+from huggingface_hub import snapshot_download
+from huggingface_hub.utils import validate_hf_hub_args
+
+
+class CustomDiffusionMixin:
+    r"""
+    Overrides DiffusionPipeline methods.
+    """
+
+    # Copied from https://github.com/huggingface/diffusers/blob/v0.30.3/src/diffusers/pipelines/pipeline_utils.py#L480
+    @classmethod
+    @validate_hf_hub_args
+    def from_pretrained(cls, pretrained_model_name_or_path, progress=None, **kwargs):
+        torch_dtype = kwargs.pop("torch_dtype", None)
+        variant = kwargs.pop("variant", None)
+        token = kwargs.pop("token", None)
+
+        # download the checkpoints and configs
+        cached_folder = cls.download(
+            pretrained_model_name_or_path,
+            variant=variant,
+            token=token,
+            **kwargs,
+        )
+
+        # pop out "_ignore_files" as it is only needed for download
+        config_dict = cls.load_config(cached_folder)
+        config_dict.pop("_ignore_files", None)
+
+        # Define which model components should load variants.
+        # We retrieve the information by matching whether variant model checkpoints exist in the subfolders.
+        # Example: `diffusion_pytorch_model.safetensors` -> `diffusion_pytorch_model.fp16.safetensors` with variant being `"fp16"`.
+        model_variants = {}
+        if variant is not None:
+            for folder in os.listdir(cached_folder):
+                folder_path = os.path.join(cached_folder, folder)
+                is_folder = os.path.isdir(folder_path) and folder in config_dict
+                variant_exists = is_folder and any(
+                    p.split(".")[1].startswith(variant) for p in os.listdir(folder_path)
+                )
+                if variant_exists:
+                    model_variants[folder] = variant
+
+        # load the pipeline class
+        pipeline_class = _get_pipeline_class(cls, config=config_dict)
+
+        # define expected modules given pipeline signature and define non-None initialized modules (=`init_kwargs`)
+        expected_modules, optional_kwargs = cls._get_signature_keys(pipeline_class)
+        passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}
+        passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs}
+
+        def load_module(name, value):
+            if value[0] is None:
+                return False
+            if name in passed_class_obj and passed_class_obj[name] is None:
+                return False
+            return True
+
+        init_dict, _, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
+        init_kwargs = {
+            k: init_dict.pop(k)
+            for k in optional_kwargs
+            if k in init_dict and k not in pipeline_class._optional_components
+        }
+        init_kwargs = {**init_kwargs, **passed_pipe_kwargs}
+        init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)}
+
+        # load each module in the pipeline
+        pipelines = import_module("diffusers.pipelines")
+        tqdm = logging.tqdm if progress is None else progress.tqdm
+        for name, (library_name, class_name) in tqdm(
+            sorted(init_dict.items()),
+            desc="Loading pipeline components",
+        ):
+            # use passed sub model or load class_name from library_name
+            loaded_sub_model = None
+            if name in passed_class_obj:
+                # passed as an argument like "scheduler"
+                loaded_sub_model = passed_class_obj[name]
+            else:
+                loaded_sub_model = load_sub_model(
+                    library_name=library_name,
+                    class_name=class_name,
+                    importable_classes=ALL_IMPORTABLE_CLASSES,
+                    pipelines=pipelines,
+                    is_pipeline_module=hasattr(pipelines, library_name),
+                    pipeline_class=pipeline_class,
+                    torch_dtype=torch_dtype,
+                    provider=None,
+                    sess_options=None,
+                    device_map=None,
+                    max_memory=None,
+                    offload_folder=None,
+                    offload_state_dict=False,
+                    model_variants=model_variants,
+                    name=name,
+                    from_flax=False,
+                    variant=variant,
+                    low_cpu_mem_usage=_LOW_CPU_MEM_USAGE_DEFAULT,
+                    cached_folder=cached_folder,
+                )
+            init_kwargs[name] = loaded_sub_model
+
+        # potentially add passed objects if expected
+        missing_modules = set(expected_modules) - set(init_kwargs.keys())
+        if len(missing_modules) > 0:
+            for module in missing_modules:
+                init_kwargs[module] = passed_class_obj.get(module, None)
+
+        # instantiate the pipeline
+        model = pipeline_class(**init_kwargs)
+
+        # save where the model was instantiated from
+        model.register_to_config(_name_or_path=pretrained_model_name_or_path)
+        return model
+
+    # Copied from https://github.com/huggingface/diffusers/blob/v0.30.3/src/diffusers/loaders/single_file.py#L270
+    @classmethod
+    @validate_hf_hub_args
+    def from_single_file(cls, pretrained_model_link_or_path, progress=None, **kwargs):
+        token = kwargs.pop("token", None)
+        torch_dtype = kwargs.pop("torch_dtype", None)
+
+        # load the pipeline class
+        pipeline_class = _get_pipeline_class(cls, config=None)
+        checkpoint = load_single_file_checkpoint(pretrained_model_link_or_path, token=token)
+
+        config = fetch_diffusers_config(checkpoint)
+        default_pretrained_model_config_name = config["pretrained_model_name_or_path"]
+
+        # attempt to download the config files for the pipeline
+        cached_model_config_path = snapshot_download(
+            default_pretrained_model_config_name,
+            token=token,
+            allow_patterns=["**/*.json", "*.json", "*.txt", "**/*.txt", "**/*.model"],
+        )
+
+        # pop out "_ignore_files" as it is only needed for download
+        config_dict = pipeline_class.load_config(cached_model_config_path)
+        config_dict.pop("_ignore_files", None)
+
+        # define expected modules given pipeline signature and define non-None initialized modules (=`init_kwargs`)
+        expected_modules, optional_kwargs = pipeline_class._get_signature_keys(cls)
+        passed_class_obj = {k: kwargs.pop(k) for k in expected_modules if k in kwargs}
+        passed_pipe_kwargs = {k: kwargs.pop(k) for k in optional_kwargs if k in kwargs}
+
+        def load_module(name, value):
+            if value[0] is None:
+                return False
+            if name in passed_class_obj and passed_class_obj[name] is None:
+                return False
+            if name in SINGLE_FILE_OPTIONAL_COMPONENTS:
+                return False
+            return True
+
+        init_dict, _, _ = pipeline_class.extract_init_dict(config_dict, **kwargs)
+        init_kwargs = {k: init_dict.pop(k) for k in optional_kwargs if k in init_dict}
+        init_kwargs = {**init_kwargs, **passed_pipe_kwargs}
+        init_dict = {k: v for k, v in init_dict.items() if load_module(k, v)}
+
+        # load each module in the pipeline
+        pipelines = import_module("diffusers.pipelines")
+        tqdm = logging.tqdm if progress is None else progress.tqdm
+        for name, (library_name, class_name) in tqdm(
+            sorted(init_dict.items()),
+            desc="Loading pipeline components",
+        ):
+            # use passed sub model or load class_name from library_name
+            loaded_sub_model = None
+            if name in passed_class_obj:
+                # passed as an argument like "scheduler"
+                loaded_sub_model = passed_class_obj[name]
+            else:
+                loaded_sub_model = load_single_file_sub_model(
+                    library_name=library_name,
+                    class_name=class_name,
+                    name=name,
+                    checkpoint=checkpoint,
+                    is_pipeline_module=hasattr(pipelines, library_name),
+                    cached_model_config_path=cached_model_config_path,
+                    pipelines=pipelines,
+                    torch_dtype=torch_dtype,
+                    **kwargs,
+                )
+            init_kwargs[name] = loaded_sub_model
+
+        # potentially add passed objects if expected
+        missing_modules = set(expected_modules) - set(init_kwargs.keys())
+        if len(missing_modules) > 0:
+            for module in missing_modules:
+                init_kwargs[module] = passed_class_obj.get(module, None)
+
+        # instantiate the pipeline
+        pipe = pipeline_class(**init_kwargs)
+
+        # save where the model was instantiated from
+        pipe.register_to_config(_name_or_path=pretrained_model_link_or_path)
+        return pipe
+
+
+class CustomStableDiffusionPipeline(CustomDiffusionMixin, StableDiffusionPipeline):
+    pass
+
+
+class CustomStableDiffusionImg2ImgPipeline(CustomDiffusionMixin, StableDiffusionImg2ImgPipeline):
+    pass

requirements.txt

@@ -10,6 +10,6 @@ httpx
 numpy==1.26.4
 peft
 ruff==0.6.7
-spaces
+spaces==0.30.2
 torch==2.2.0
 torchvision==0.17.0