Simplify loading and inference

app.py

@@ -1,11 +1,20 @@
 import argparse
+import os
+from importlib.util import find_spec
+
+# Improved GPU handling and progress bars
+os.environ["ZEROGPU_V2"] = "1"
+
+# Use Rust-based downloader
+if find_spec("hf_transfer"):
+    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 import gradio as gr
+from huggingface_hub._snapshot_download import snapshot_download
 
 from lib import (
     Config,
-    # disable_progress_bars,
-    download_repo_files,
+    disable_progress_bars,
     generate,
     read_file,
     read_json,
@@ -60,24 +69,6 @@ random_prompt_js = f"""
 }}
 """
 
-
-# Transform the raw inputs before generation
-def generate_fn(*args, progress=gr.Progress(track_tqdm=True)):
-    if len(args) > 0:
-        prompt = args[0]
-    else:
-        prompt = None
-    if prompt is None or prompt.strip() == "":
-        raise gr.Error("You must enter a prompt")
-    try:
-        # if Config.ZERO_GPU:
-        #     progress((0, 100), desc="ZeroGPU init")
-        images = generate(*args, Error=gr.Error, Info=gr.Info, progress=progress)
-    except RuntimeError:
-        raise gr.Error("Error: Please try again")
-    return images
-
-
 with gr.Blocks(
     head=read_file("./partials/head.html"),
     css="./app.css",
@@ -244,10 +235,10 @@ with gr.Blocks(
                     label="Scale",
                 )
                 seed = gr.Number(
-                    value=Config.SEED,
-                    label="Seed",
                     minimum=-1,
                     maximum=(2**64) - 1,
+                    label="Seed",
+                    value=-1,
                 )
             with gr.Row():
                 use_karras = gr.Checkbox(
@@ -293,7 +284,7 @@ with gr.Blocks(
     # Generate images
     gr.on(
         triggers=[generate_btn.click, prompt.submit],
-        fn=generate_fn,
+        fn=generate,
         api_name="generate",
         outputs=[output_images],
         inputs=[
@@ -321,8 +312,16 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     # disable_progress_bars()
+    token = os.environ.get("HF_TOKEN", None)
     for repo_id, allow_patterns in Config.HF_REPOS.items():
-        download_repo_files(repo_id, allow_patterns, token=Config.HF_TOKEN)
+        snapshot_download(
+            repo_id=repo_id,
+            repo_type="model",
+            revision="main",
+            token=token,
+            allow_patterns=allow_patterns,
+            ignore_patterns=None,
+        )
 
     # https://www.gradio.app/docs/gradio/interface#interface-queue
     demo.queue(default_concurrency_limit=1).launch(

lib/__init__.py

@@ -2,7 +2,6 @@ from .config import Config
 from .inference import generate
 from .utils import (
     disable_progress_bars,
-    download_repo_files,
     read_file,
     read_json,
 )
@@ -10,7 +9,6 @@ from .utils import (
 __all__ = [
     "Config",
     "disable_progress_bars",
-    "download_repo_files",
     "generate",
     "read_file",
     "read_json",

lib/config.py

@@ -1,6 +1,3 @@
-import os
-from importlib import import_module
-from importlib.util import find_spec
 from types import SimpleNamespace
 from warnings import filterwarnings
 
@@ -16,13 +13,6 @@ from diffusers import (
 from diffusers.utils import logging as diffusers_logging
 from transformers import logging as transformers_logging
 
-# Improved GPU handling and progress bars; set before importing spaces
-os.environ["ZEROGPU_V2"] = "1"
-
-# Use Rust-based downloader; errors if enabled and not installed
-if find_spec("hf_transfer"):
-    os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
-
 filterwarnings("ignore", category=FutureWarning, module="diffusers")
 filterwarnings("ignore", category=FutureWarning, module="transformers")
 
@@ -60,20 +50,10 @@ _sdxl_files_with_vae = [*_sdxl_files, "vae_1_0/config.json"]
 
 # Using namespace instead of dataclass for simplicity
 Config = SimpleNamespace(
-    HF_TOKEN=os.environ.get("HF_TOKEN", None),
-    ZERO_GPU=import_module("spaces").config.Config.zero_gpu,
     PIPELINES={
         "txt2img": StableDiffusionXLPipeline,
         "img2img": StableDiffusionXLImg2ImgPipeline,
     },
-    MODEL="segmind/Segmind-Vega",
-    MODELS=[
-        "cyberdelia/CyberRealsticXL",
-        "fluently/Fluently-XL-Final",
-        "segmind/Segmind-Vega",
-        "SG161222/RealVisXL_V5.0",
-        "stabilityai/stable-diffusion-xl-base-1.0",
-    ],
     HF_REPOS={
         "ai-forever/Real-ESRGAN": ["RealESRGAN_x2.pth", "RealESRGAN_x4.pth"],
         "cyberdelia/CyberRealsticXL": ["CyberRealisticXLPlay_V1.0.safetensors"],
@@ -84,10 +64,18 @@ Config = SimpleNamespace(
         "stabilityai/stable-diffusion-xl-base-1.0": _sdxl_files_with_vae,
         "stabilityai/stable-diffusion-xl-refiner-1.0": _sdxl_refiner_files,
     },
+    MODEL="segmind/Segmind-Vega",
+    MODELS=[
+        "cyberdelia/CyberRealsticXL",
+        "fluently/Fluently-XL-Final",
+        "segmind/Segmind-Vega",
+        "SG161222/RealVisXL_V5.0",
+        "stabilityai/stable-diffusion-xl-base-1.0",
+    ],
     SINGLE_FILE_MODELS=[
-        "cyberdelia/cyberrealsticxl",
-        "fluently/fluently-xl-final",
-        "sg161222/realvisxl_v5.0",
+        "cyberdelia/CyberRealsticXL",
+        "fluently/Fluently-XL-Final",
+        "SG161222/RealVisXL_V5.0",
     ],
     VAE_MODEL="madebyollin/sdxl-vae-fp16-fix",
     REFINER_MODEL="stabilityai/stable-diffusion-xl-refiner-1.0",
@@ -102,7 +90,6 @@ Config = SimpleNamespace(
     WIDTH=1024,
     HEIGHT=1024,
     NUM_IMAGES=1,
-    SEED=-1,
     GUIDANCE_SCALE=6,
     INFERENCE_STEPS=40,
     DEEPCACHE_INTERVAL=1,

lib/inference.py

@@ -4,40 +4,17 @@ from datetime import datetime
 import torch
 from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import PromptParser
+from gradio import Error, Info, Progress
 from spaces import GPU
 
-from .config import Config
 from .loader import Loader
 from .logger import Logger
-from .utils import cuda_collect, safe_progress, timer
-
-
-# Dynamic signature for the GPU duration function; max 60s per image
-def gpu_duration(**kwargs):
-    loading = 15
-    duration = 15
-    width = kwargs.get("width", 1024)
-    height = kwargs.get("height", 1024)
-    scale = kwargs.get("scale", 1)
-    num_images = kwargs.get("num_images", 1)
-    use_refiner = kwargs.get("use_refiner", False)
-    size = width * height
-    if use_refiner:
-        loading += 10
-    if size > 1_100_000:
-        duration += 5
-    if size > 1_600_000:
-        duration += 5
-    if scale == 2:
-        duration += 5
-    if scale == 4:
-        duration += 10
-    return loading + (duration * num_images)
-
-
-@GPU(duration=gpu_duration)
+from .utils import cuda_collect, get_output_types, timer
+
+
+@GPU
 def generate(
-    positive_prompt,
+    positive_prompt="",
     negative_prompt="",
     seed=None,
     model="stabilityai/stable-diffusion-xl-base-1.0",
@@ -51,50 +28,21 @@ def generate(
     num_images=1,
     use_karras=False,
     use_refiner=False,
-    Error=Exception,
-    Info=None,
-    progress=None,
+    progress=Progress(track_tqdm=True),
 ):
+    if not torch.cuda.is_available():
+        raise Error("CUDA not available")
+
+    if positive_prompt.strip() == "":
+        raise Error("You must enter a prompt")
+
     KIND = "txt2img"
-    CURRENT_STEP = 0
-    CURRENT_IMAGE = 1
     EMBEDDINGS_TYPE = ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED
 
     start = time.perf_counter()
     log = Logger("generate")
     log.info(f"Generating {num_images} image{'s' if num_images > 1 else ''}...")
 
-    if Config.ZERO_GPU:
-        safe_progress(progress, 100, 100, "ZeroGPU init")
-
-    if not torch.cuda.is_available():
-        raise Error("CUDA not available")
-
-    # https://pytorch.org/docs/stable/generated/torch.manual_seed.html
-    if seed is None or seed < 0:
-        seed = int(datetime.now().timestamp() * 1e6) % (2**64)
-
-    # custom progress bar for multiple images
-    def callback_on_step_end(pipeline, step, timestep, latents):
-        nonlocal CURRENT_IMAGE, CURRENT_STEP
-        if progress is not None:
-            # calculate total steps for img2img based on denoising strength
-            strength = 1
-            total_steps = min(int(inference_steps * strength), inference_steps)
-
-            # if steps are different we're in the refiner
-            refining = False
-            if CURRENT_STEP == step:
-                CURRENT_STEP = step + 1
-            else:
-                refining = True
-                CURRENT_STEP += 1
-            progress(
-                (CURRENT_STEP, total_steps),
-                desc=f"{'Refining' if refining else 'Generating'} image {CURRENT_IMAGE}/{num_images}",
-            )
-        return latents
-
     loader = Loader()
     loader.load(
         KIND,
@@ -111,10 +59,11 @@ def generate(
     pipeline = loader.pipeline
     upscaler = loader.upscaler
 
+    # Probably a typo in the config
     if pipeline is None:
         raise Error(f"Error loading {model}")
 
-    # prompt embeds for base and refiner
+    # Prompt embeddings for base and refiner
     compel_1 = Compel(
         text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
         tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
@@ -132,9 +81,13 @@ def generate(
         device=pipeline.device,
     )
 
+    # https://pytorch.org/docs/stable/generated/torch.manual_seed.html
+    if seed is None or seed < 0:
+        seed = int(datetime.now().timestamp() * 1e6) % (2**64)
+
+    # Increment the seed after each iteration
     images = []
     current_seed = seed
-    safe_progress(progress, 0, num_images, f"Generating image 0/{num_images}")
 
     for i in range(num_images):
         try:
@@ -144,23 +97,14 @@ def generate(
         except PromptParser.ParsingException:
             raise Error("Invalid prompt")
 
-        # refiner expects latents; upscaler expects numpy array
-        pipe_output_type = "pil"
-        refiner_output_type = "pil"
-        if use_refiner:
-            pipe_output_type = "latent"
-            if scale > 1:
-                refiner_output_type = "np"
-        else:
-            if scale > 1:
-                pipe_output_type = "np"
-
-        pipe_kwargs = {
+        pipeline_output_type, refiner_output_type = get_output_types(scale, use_refiner)
+
+        pipeline_kwargs = {
             "width": width,
             "height": height,
             "denoising_end": 0.8 if use_refiner else None,
             "generator": generator,
-            "output_type": pipe_output_type,
+            "output_type": pipeline_output_type,
             "guidance_scale": guidance_scale,
             "num_inference_steps": inference_steps,
             "prompt_embeds": conditioning_1[0:1],
@@ -181,39 +125,31 @@ def generate(
             "negative_pooled_prompt_embeds": pooled_2[1:2],
         }
 
-        if progress is not None:
-            pipe_kwargs["callback_on_step_end"] = callback_on_step_end
-            refiner_kwargs["callback_on_step_end"] = callback_on_step_end
+        image = pipeline(**pipeline_kwargs).images[0]
 
-        try:
-            image = pipeline(**pipe_kwargs).images[0]
-            if use_refiner:
-                refiner_kwargs["image"] = image
-                image = refiner(**refiner_kwargs).images[0]
-            images.append((image, str(current_seed)))
-            current_seed += 1
-        finally:
-            CURRENT_STEP = 0
-            CURRENT_IMAGE += 1
+        if use_refiner:
+            refiner_kwargs["image"] = image
+            image = refiner(**refiner_kwargs).images[0]
+
+        # Use a tuple so gallery images get captions
+        images.append((image, str(current_seed)))
+        current_seed += 1
 
     # Upscale
     if scale > 1:
-        msg = f"Upscaling {scale}x"
-        with timer(msg):
-            safe_progress(progress, 0, num_images, desc=msg)
+        with timer(f"Upscaling {num_images} images {scale}x", logger=log.info):
             for i, image in enumerate(images):
-                images = upscaler.predict(image[0])
-                images[i] = image
-                safe_progress(progress, i + 1, num_images, desc=msg)
+                image = upscaler.predict(image[0])
+                seed = images[i][1]
+                images[i] = (image, seed)
 
-    # Flush memory after generating
+    # Flush cache after generating
     cuda_collect()
 
     end = time.perf_counter()
     msg = f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {end - start:.2f}s"
     log.info(msg)
 
-    # Alert if notifier provided
     if Info:
         Info(msg)
 

lib/loader.py

@@ -1,5 +1,3 @@
-# import gc
-
 import torch
 from DeepCache import DeepCacheSDHelper
 from diffusers.models import AutoencoderKL
@@ -33,7 +31,7 @@ class Loader:
         return False
 
     def should_unload_pipeline(self, model=""):
-        return self.pipeline is not None and self.model.lower() != model.lower()
+        return self.pipeline is not None and self.model != model
 
     def should_load_refiner(self, use_refiner=False):
         return self.refiner is None and use_refiner
@@ -53,8 +51,6 @@ class Loader:
         return self.pipeline is None
 
     def unload(self, model, use_refiner, deepcache_interval, scale):
-        needs_gc = False
-
         if self.should_unload_deepcache(deepcache_interval):
             self.log.info("Disabling DeepCache")
             self.pipeline.deepcache.disable()
@@ -64,37 +60,41 @@ class Loader:
                 delattr(self.refiner, "deepcache")
 
         if self.should_unload_refiner(use_refiner):
-            with timer("Unloading refiner"):
-                self.refiner.to("cpu", silence_dtype_warnings=True)
-                self.refiner = None
-                needs_gc = True
+            self.log.info("Unloading refiner")
+            self.refiner = None
 
         if self.should_unload_upscaler(scale):
-            with timer(f"Unloading {self.upscaler.scale}x upscaler"):
-                self.upscaler.to("cpu")
-                self.upscaler = None
-                needs_gc = True
+            self.log.info("Unloading upscaler")
+            self.upscaler = None
 
         if self.should_unload_pipeline(model):
-            with timer(f"Unloading {self.model}"):
-                self.pipeline.to("cpu", silence_dtype_warnings=True)
-                if self.refiner:
-                    self.refiner.vae = None
-                    self.refiner.scheduler = None
-                    self.refiner.tokenizer_2 = None
-                    self.refiner.text_encoder_2 = None
-                self.pipeline = None
-                self.model = None
-                needs_gc = True
-
-        if needs_gc:
-            cuda_collect()
-            # gc.collect()
-
-    def load_refiner(self, refiner_kwargs={}, progress=None):
+            self.log.info(f"Unloading {self.model}")
+            if self.refiner:
+                self.refiner.vae = None
+                self.refiner.scheduler = None
+                self.refiner.tokenizer_2 = None
+                self.refiner.text_encoder_2 = None
+            self.pipeline = None
+            self.model = None
+
+        # Flush cache
+        cuda_collect()
+
+    def load_refiner(self, progress=None):
         model = Config.REFINER_MODEL
         try:
-            with timer(f"Loading {model}"):
+            with timer(f"Loading {model}", logger=self.log.info):
+                refiner_kwargs = {
+                    "variant": "fp16",
+                    "torch_dtype": self.pipeline.dtype,
+                    "add_watermarker": False,
+                    "requires_aesthetics_score": True,
+                    "force_zeros_for_empty_prompt": False,
+                    "vae": self.pipeline.vae,
+                    "scheduler": self.pipeline.scheduler,
+                    "tokenizer_2": self.pipeline.tokenizer_2,
+                    "text_encoder_2": self.pipeline.text_encoder_2,
+                }
                 Pipeline = Config.PIPELINES["img2img"]
                 self.refiner = Pipeline.from_pretrained(model, **refiner_kwargs).to("cuda")
         except Exception as e:
@@ -107,7 +107,7 @@ class Loader:
     def load_upscaler(self, scale=1):
         if self.should_load_upscaler(scale):
             try:
-                with timer(f"Loading {scale}x upscaler"):
+                with timer(f"Loading {scale}x upscaler", logger=self.log.info):
                     self.upscaler = RealESRGAN(scale, device=self.pipeline.device)
                     self.upscaler.load_weights()
             except Exception as e:
@@ -125,7 +125,7 @@ class Loader:
                 self.refiner.deepcache.set_params(cache_interval=interval)
                 self.refiner.deepcache.enable()
 
-    def load(self, kind, model, scheduler, deepcache_interval, scale, use_karras, use_refiner, progress):
+    def load(self, kind, model, scheduler, deepcache_interval, scale, use_karras, use_refiner, progress=None):
         scheduler_kwargs = {
             "beta_start": 0.00085,
             "beta_end": 0.012,
@@ -141,13 +141,13 @@ class Loader:
             scheduler_kwargs["clip_sample"] = False
             scheduler_kwargs["set_alpha_to_one"] = False
 
-        if model.lower() not in Config.SINGLE_FILE_MODELS:
+        if model not in Config.SINGLE_FILE_MODELS:
             variant = "fp16"
         else:
             variant = None
 
         dtype = torch.float16
-        pipe_kwargs = {
+        pipeline_kwargs = {
             "variant": variant,
             "torch_dtype": dtype,
             "add_watermarker": False,
@@ -161,16 +161,16 @@ class Loader:
         Scheduler = Config.SCHEDULERS[scheduler]
 
         try:
-            with timer(f"Loading {model}"):
+            with timer(f"Loading {model}", logger=self.log.info):
                 self.model = model
-                if model.lower() in Config.SINGLE_FILE_MODELS:
+                if model in Config.SINGLE_FILE_MODELS:
                     checkpoint = Config.HF_REPOS[model][0]
                     self.pipeline = Pipeline.from_single_file(
                         f"https://huggingface.co/{model}/{checkpoint}",
-                        **pipe_kwargs,
+                        **pipeline_kwargs,
                     ).to("cuda")
                 else:
-                    self.pipeline = Pipeline.from_pretrained(model, **pipe_kwargs).to("cuda")
+                    self.pipeline = Pipeline.from_pretrained(model, **pipeline_kwargs).to("cuda")
         except Exception as e:
             self.log.error(f"Error loading {model}: {e}")
             self.model = None
@@ -190,7 +190,7 @@ class Loader:
             or self.pipeline.scheduler.config.use_karras_sigmas == use_karras
         )
 
-        if self.model.lower() == model.lower():
+        if self.model == model:
             if not same_scheduler:
                 self.log.info(f"Enabling {scheduler}")
             if not same_karras:
@@ -201,18 +201,7 @@ class Loader:
                     self.refiner.scheduler = self.pipeline.scheduler
 
         if self.should_load_refiner(use_refiner):
-            refiner_kwargs = {
-                "variant": "fp16",
-                "torch_dtype": dtype,
-                "add_watermarker": False,
-                "requires_aesthetics_score": True,
-                "force_zeros_for_empty_prompt": False,
-                "vae": self.pipeline.vae,
-                "scheduler": self.pipeline.scheduler,
-                "tokenizer_2": self.pipeline.tokenizer_2,
-                "text_encoder_2": self.pipeline.text_encoder_2,
-            }
-            self.load_refiner(refiner_kwargs, progress)
+            self.load_refiner(progress)
 
         if self.should_load_deepcache(deepcache_interval):
             self.load_deepcache(deepcache_interval)

lib/utils.py

@@ -5,8 +5,6 @@ from contextlib import contextmanager
 
 import torch
 from diffusers.utils import logging as diffusers_logging
-from huggingface_hub._snapshot_download import snapshot_download
-from huggingface_hub.utils import are_progress_bars_disabled
 from transformers import logging as transformers_logging
 
 
@@ -45,9 +43,14 @@ def enable_progress_bars():
     diffusers_logging.enable_progress_bar()
 
 
-def safe_progress(progress, current=0, total=0, desc=""):
-    if progress is not None:
-        progress((current, total), desc=desc)
+def get_output_types(scale=1, use_refiner=False):
+    if use_refiner:
+        pipeline_type = "latent"
+        refiner_type = "np" if scale > 1 else "pil"
+    else:
+        refiner_type = "pil"
+        pipeline_type = "np" if scale > 1 else "pil"
+    return (pipeline_type, refiner_type)
 
 
 def cuda_collect():
@@ -56,19 +59,3 @@ def cuda_collect():
         torch.cuda.ipc_collect()
         torch.cuda.reset_peak_memory_stats()
         torch.cuda.synchronize()
-
-
-def download_repo_files(repo_id, allow_patterns, token=None):
-    was_disabled = are_progress_bars_disabled()
-    enable_progress_bars()
-    snapshot_path = snapshot_download(
-        repo_id=repo_id,
-        repo_type="model",
-        revision="main",
-        token=token,
-        allow_patterns=allow_patterns,
-        ignore_patterns=None,
-    )
-    if was_disabled:
-        disable_progress_bars()
-    return snapshot_path