Clean up loader

.gitignore

@@ -1,2 +1,3 @@
 __pycache__/
 .venv/
+app.log

app.py

@@ -2,7 +2,14 @@ import argparse
 
 import gradio as gr
 
-from lib import Config, async_call, disable_progress_bars, download_repo_files, generate, read_file, read_json
+from lib import (
+    Config,
+    # disable_progress_bars,
+    download_repo_files,
+    generate,
+    read_file,
+    read_json,
+)
 
 # Update refresh button hover text
 seed_js = """
@@ -55,28 +62,19 @@ random_prompt_js = f"""
 
 
 # Transform the raw inputs before generation
-async def generate_fn(*args, progress=gr.Progress(track_tqdm=True)):
+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 = await async_call(
-            generate,
-            *args,
-            Error=gr.Error,
-            Info=gr.Info,
-            progress=progress,
-        )
+        # 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
 
 
@@ -259,7 +257,7 @@ with gr.Blocks(
                 )
                 use_refiner = gr.Checkbox(
                     elem_classes=["checkbox"],
-                    label="Refiner",
+                    label="Use refiner",
                     value=False,
                 )
 
@@ -322,8 +320,8 @@ if __name__ == "__main__":
     parser.add_argument("-p", "--port", type=int, metavar="INT", default=7860)
     args = parser.parse_args()
 
-    disable_progress_bars()
-    for repo_id, allow_patterns in Config.HF_MODELS.items():
+    # disable_progress_bars()
+    for repo_id, allow_patterns in Config.HF_REPOS.items():
         download_repo_files(repo_id, allow_patterns, token=Config.HF_TOKEN)
 
     # https://www.gradio.app/docs/gradio/interface#interface-queue

lib/__init__.py

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

lib/config.py

@@ -56,15 +56,12 @@ _sdxl_files = [
     "tokenizer/vocab.json",
 ]
 
+_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,
-    HF_MODELS={
-        "segmind/Segmind-Vega": [*_sdxl_files],
-        "stabilityai/stable-diffusion-xl-base-1.0": [*_sdxl_files, "vae_1_0/config.json"],
-        "stabilityai/stable-diffusion-xl-refiner-1.0": [*_sdxl_refiner_files],
-    },
     PIPELINES={
         "txt2img": StableDiffusionXLPipeline,
         "img2img": StableDiffusionXLImg2ImgPipeline,
@@ -77,13 +74,21 @@ Config = SimpleNamespace(
         "SG161222/RealVisXL_V5.0",
         "stabilityai/stable-diffusion-xl-base-1.0",
     ],
-    # Single-file model weights
-    MODEL_CHECKPOINTS={
-        # keep keys lowercase for case-insensitive matching in the loader
-        "cyberdelia/cyberrealsticxl": "CyberRealisticXLPlay_V1.0.safetensors",  # typo in "realistic"
-        "fluently/fluently-xl-final": "FluentlyXL-Final.safetensors",
-        "sg161222/realvisxl_v5.0": "RealVisXL_V5.0_fp16.safetensors",
+    HF_REPOS={
+        "ai-forever/Real-ESRGAN": ["RealESRGAN_x2.pth", "RealESRGAN_x4.pth"],
+        "cyberdelia/CyberRealsticXL": ["CyberRealisticXLPlay_V1.0.safetensors"],
+        "fluently/Fluently-XL-Final": ["FluentlyXL-Final.safetensors"],
+        "madebyollin/sdxl-vae-fp16-fix": ["config.json", "diffusion_pytorch_model.fp16.safetensors"],
+        "segmind/Segmind-Vega": _sdxl_files,
+        "SG161222/RealVisXL_V5.0": ["RealVisXL_V5.0_fp16.safetensors"],
+        "stabilityai/stable-diffusion-xl-base-1.0": _sdxl_files_with_vae,
+        "stabilityai/stable-diffusion-xl-refiner-1.0": _sdxl_refiner_files,
     },
+    SINGLE_FILE_MODELS=[
+        "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",
     SCHEDULER="Euler",

lib/inference.py

@@ -9,7 +9,7 @@ from spaces import GPU
 from .config import Config
 from .loader import Loader
 from .logger import Logger
-from .utils import clear_cuda_cache, safe_progress, timer
+from .utils import cuda_collect, safe_progress, timer
 
 
 # Dynamic signature for the GPU duration function; max 60s per image
@@ -55,6 +55,11 @@ def generate(
     Info=None,
     progress=None,
 ):
+    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 ''}...")
@@ -69,11 +74,6 @@ def generate(
     if seed is None or seed < 0:
         seed = int(datetime.now().timestamp() * 1e6) % (2**64)
 
-    KIND = "txt2img"
-    CURRENT_STEP = 0
-    CURRENT_IMAGE = 1
-    EMBEDDINGS_TYPE = ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED
-
     # custom progress bar for multiple images
     def callback_on_step_end(pipeline, step, timestep, latents):
         nonlocal CURRENT_IMAGE, CURRENT_STEP
@@ -107,29 +107,29 @@ def generate(
         progress,
     )
 
-    if loader.pipe is None:
-        raise Error(f"Error loading {model}")
-
-    pipe = loader.pipe
     refiner = loader.refiner
+    pipeline = loader.pipeline
     upscaler = loader.upscaler
 
+    if pipeline is None:
+        raise Error(f"Error loading {model}")
+
     # prompt embeds for base and refiner
     compel_1 = Compel(
-        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
-        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+        text_encoder=[pipeline.text_encoder, pipeline.text_encoder_2],
+        tokenizer=[pipeline.tokenizer, pipeline.tokenizer_2],
         requires_pooled=[False, True],
         returned_embeddings_type=EMBEDDINGS_TYPE,
-        dtype_for_device_getter=lambda _: pipe.dtype,
-        device=pipe.device,
+        dtype_for_device_getter=lambda _: pipeline.dtype,
+        device=pipeline.device,
     )
     compel_2 = Compel(
-        text_encoder=[pipe.text_encoder_2],
-        tokenizer=[pipe.tokenizer_2],
+        text_encoder=[pipeline.text_encoder_2],
+        tokenizer=[pipeline.tokenizer_2],
         requires_pooled=[True],
         returned_embeddings_type=EMBEDDINGS_TYPE,
-        dtype_for_device_getter=lambda _: pipe.dtype,
-        device=pipe.device,
+        dtype_for_device_getter=lambda _: pipeline.dtype,
+        device=pipeline.device,
     )
 
     images = []
@@ -138,7 +138,7 @@ def generate(
 
     for i in range(num_images):
         try:
-            generator = torch.Generator(device=pipe.device).manual_seed(current_seed)
+            generator = torch.Generator(device=pipeline.device).manual_seed(current_seed)
             conditioning_1, pooled_1 = compel_1([positive_prompt, negative_prompt])
             conditioning_2, pooled_2 = compel_2([positive_prompt, negative_prompt])
         except PromptParser.ParsingException:
@@ -186,7 +186,7 @@ def generate(
             refiner_kwargs["callback_on_step_end"] = callback_on_step_end
 
         try:
-            image = pipe(**pipe_kwargs).images[0]
+            image = pipeline(**pipe_kwargs).images[0]
             if use_refiner:
                 refiner_kwargs["image"] = image
                 image = refiner(**refiner_kwargs).images[0]
@@ -207,7 +207,7 @@ def generate(
                 safe_progress(progress, i + 1, num_images, desc=msg)
 
     # Flush memory after generating
-    clear_cuda_cache()
+    cuda_collect()
 
     end = time.perf_counter()
     msg = f"Generated {len(images)} image{'s' if len(images) > 1 else ''} in {end - start:.2f}s"

lib/loader.py

@@ -1,5 +1,4 @@
-import gc
-from threading import Lock
+# import gc
 
 import torch
 from DeepCache import DeepCacheSDHelper
@@ -8,194 +7,125 @@ from diffusers.models import AutoencoderKL
 from .config import Config
 from .logger import Logger
 from .upscaler import RealESRGAN
-from .utils import clear_cuda_cache, timer
+from .utils import cuda_collect, timer
 
 
 class Loader:
-    _instance = None
-    _lock = Lock()
-
-    def __new__(cls):
-        with cls._lock:
-            if cls._instance is None:
-                cls._instance = super().__new__(cls)
-                cls._instance.pipe = None
-                cls._instance.model = None
-                cls._instance.refiner = None
-                cls._instance.upscaler = None
-                cls._instance.log = Logger("Loader")
-        return cls._instance
-
-    def _should_unload_refiner(self, refiner=False):
-        if self.refiner is None:
-            return False
-        if not refiner:
-            return True
-        return False
+    def __init__(self):
+        self.model = ""
+        self.refiner = None
+        self.pipeline = None
+        self.upscaler = None
+        self.log = Logger("Loader")
 
-    def _should_unload_upscaler(self, scale=1):
-        if self.upscaler is not None and self.upscaler.scale != scale:
-            return True
-        return False
+    def should_unload_refiner(self, use_refiner=False):
+        return self.refiner is not None and not use_refiner
 
-    def _should_unload_deepcache(self, interval=1):
-        has_deepcache = hasattr(self.pipe, "deepcache")
+    def should_unload_upscaler(self, scale=1):
+        return self.upscaler is not None and self.upscaler.scale != scale
+
+    def should_unload_deepcache(self, interval=1):
+        has_deepcache = hasattr(self.pipeline, "deepcache")
         if has_deepcache and interval == 1:
             return True
-        if has_deepcache and self.pipe.deepcache.params["cache_interval"] != interval:
+        if has_deepcache and self.pipeline.deepcache.params["cache_interval"] != interval:
             return True
         return False
 
-    def _should_unload_pipeline(self, model=""):
-        if self.pipe is None:
-            return False
-        if self.model and self.model.lower() != model.lower():
+    def should_unload_pipeline(self, model=""):
+        return self.pipeline is not None and self.model.lower() != model.lower()
+
+    def should_load_refiner(self, use_refiner=False):
+        return self.refiner is None and use_refiner
+
+    def should_load_upscaler(self, scale=1):
+        return self.upscaler is None and scale > 1
+
+    def should_load_deepcache(self, interval=1):
+        has_deepcache = hasattr(self.pipeline, "deepcache")
+        if not has_deepcache and interval != 1:
+            return True
+        if has_deepcache and self.pipeline.deepcache.params["cache_interval"] != interval:
             return True
         return False
 
-    def _unload_refiner(self):
-        if self.refiner is not None:
+    def should_load_pipeline(self):
+        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()
+            delattr(self.pipeline, "deepcache")
+            if self.refiner:
+                self.refiner.deepcache.disable()
+                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
 
-    def _unload_upscaler(self):
-        if self.upscaler is not 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
 
-    def _unload_deepcache(self):
-        if self.pipe.deepcache is not None:
-            self.log.info("Disabling DeepCache")
-            self.pipe.deepcache.disable()
-            delattr(self.pipe, "deepcache")
-            if self.refiner is not None:
-                if hasattr(self.refiner, "deepcache"):
-                    self.refiner.deepcache.disable()
-                    delattr(self.refiner, "deepcache")
-
-    def _unload_pipeline(self):
-        if self.pipe is not None:
+        if self.should_unload_pipeline(model):
             with timer(f"Unloading {self.model}"):
-                self.pipe.to("cpu", silence_dtype_warnings=True)
-                if self.refiner is not None:
+                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
-
-    def _unload(self, model, refiner, deepcache, scale):
-        to_unload = []
-        if self._should_unload_deepcache(deepcache):  # remove deepcache first
-            self._unload_deepcache()
-
-        if self._should_unload_refiner(refiner):
-            self._unload_refiner()
-            to_unload.append("refiner")
-
-        if self._should_unload_upscaler(scale):
-            self._unload_upscaler()
-            to_unload.append("upscaler")
-
-        if self._should_unload_pipeline(model):
-            self._unload_pipeline()
-            to_unload.append("model")
-            to_unload.append("pipe")
-
-        # Flush cache and run garbage collector
-        clear_cuda_cache()
-        for component in to_unload:
-            setattr(self, component, None)
-        gc.collect()
-
-    def _should_load_refiner(self, refiner=False):
-        if self.refiner is None and refiner:
-            return True
-        return False
-
-    def _should_load_upscaler(self, scale=1):
-        if self.upscaler is None and scale > 1:
-            return True
-        return False
-
-    def _should_load_deepcache(self, interval=1):
-        has_deepcache = hasattr(self.pipe, "deepcache")
-        if not has_deepcache and interval != 1:
-            return True
-        if has_deepcache and self.pipe.deepcache.params["cache_interval"] != interval:
-            return True
-        return False
-
-    def _should_load_pipeline(self):
-        if self.pipe is None:
-            return True
-        return False
-
-    def _load_refiner(self, refiner, progress, **kwargs):
-        if self._should_load_refiner(refiner):
-            model = Config.REFINER_MODEL
-            pipeline = Config.PIPELINES["img2img"]
-            try:
-                with timer(f"Loading {model}"):
-                    self.refiner = pipeline.from_pretrained(model, **kwargs).to("cuda")
-            except Exception as e:
-                self.log.error(f"Error loading {model}: {e}")
-                self.refiner = None
-                return
+                self.pipeline = None
+                self.model = None
+                needs_gc = True
+
+        if needs_gc:
+            cuda_collect()
+            # gc.collect()
+
+    def load_refiner(self, refiner_kwargs={}, progress=None):
+        model = Config.REFINER_MODEL
+        try:
+            with timer(f"Loading {model}"):
+                Pipeline = Config.PIPELINES["img2img"]
+                self.refiner = Pipeline.from_pretrained(model, **refiner_kwargs).to("cuda")
+        except Exception as e:
+            self.log.error(f"Error loading {model}: {e}")
+            self.refiner = None
+            return
         if self.refiner is not None:
             self.refiner.set_progress_bar_config(disable=progress is not None)
 
-    def _load_upscaler(self, scale=1):
-        if self._should_load_upscaler(scale):
+    def load_upscaler(self, scale=1):
+        if self.should_load_upscaler(scale):
             try:
                 with timer(f"Loading {scale}x upscaler"):
-                    self.upscaler = RealESRGAN(scale, device=self.pipe.device)
+                    self.upscaler = RealESRGAN(scale, device=self.pipeline.device)
                     self.upscaler.load_weights()
             except Exception as e:
                 self.log.error(f"Error loading {scale}x upscaler: {e}")
                 self.upscaler = None
 
-    def _load_deepcache(self, interval=1):
-        if self._should_load_deepcache(interval):
+    def load_deepcache(self, interval=1):
+        if self.should_load_deepcache(interval):
             self.log.info("Enabling DeepCache")
-            self.pipe.deepcache = DeepCacheSDHelper(pipe=self.pipe)
-            self.pipe.deepcache.set_params(cache_interval=interval)
-            self.pipe.deepcache.enable()
-            if self.refiner is not None:
+            self.pipeline.deepcache = DeepCacheSDHelper(pipe=self.pipeline)
+            self.pipeline.deepcache.set_params(cache_interval=interval)
+            self.pipeline.deepcache.enable()
+            if self.refiner:
                 self.refiner.deepcache = DeepCacheSDHelper(pipe=self.refiner)
                 self.refiner.deepcache.set_params(cache_interval=interval)
                 self.refiner.deepcache.enable()
 
-    def _load_pipeline(self, kind, model, progress, **kwargs):
-        pipeline = Config.PIPELINES[kind]
-        if self._should_load_pipeline():
-            try:
-                with timer(f"Loading {model}"):
-                    self.model = model
-                    if model.lower() in Config.MODEL_CHECKPOINTS.keys():
-                        self.pipe = pipeline.from_single_file(
-                            f"https://huggingface.co/{model}/{Config.MODEL_CHECKPOINTS[model.lower()]}",
-                            **kwargs,
-                        ).to("cuda")
-                    else:
-                        self.pipe = pipeline.from_pretrained(model, **kwargs).to("cuda")
-                    if self.refiner is not None:
-                        self.refiner.vae = self.pipe.vae
-                        self.refiner.scheduler = self.pipe.scheduler
-                        self.refiner.tokenizer_2 = self.pipe.tokenizer_2
-                        self.refiner.text_encoder_2 = self.pipe.text_encoder_2
-                        self.refiner.to(self.pipe.device)
-            except Exception as e:
-                self.log.error(f"Error loading {model}: {e}")
-                self.model = None
-                self.pipe = None
-                self.refiner = None
-                return
-        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=progress is not None)
-
-    def load(self, kind, model, scheduler, deepcache, scale, karras, refiner, progress):
+    def load(self, kind, model, scheduler, deepcache_interval, scale, use_karras, use_refiner, progress):
         scheduler_kwargs = {
             "beta_start": 0.00085,
             "beta_end": 0.012,
@@ -205,14 +135,13 @@ class Loader:
         }
 
         if scheduler not in ["DDIM", "Euler a"]:
-            scheduler_kwargs["use_karras_sigmas"] = karras
+            scheduler_kwargs["use_karras_sigmas"] = use_karras
 
-        # https://github.com/huggingface/diffusers/blob/8a3f0c1/scripts/convert_original_stable_diffusion_to_diffusers.py#L939
         if scheduler == "DDIM":
             scheduler_kwargs["clip_sample"] = False
             scheduler_kwargs["set_alpha_to_one"] = False
 
-        if model.lower() not in Config.MODEL_CHECKPOINTS.keys():
+        if model.lower() not in Config.SINGLE_FILE_MODELS:
             variant = "fp16"
         else:
             variant = None
@@ -226,47 +155,67 @@ class Loader:
             "vae": AutoencoderKL.from_pretrained(Config.VAE_MODEL, torch_dtype=dtype),
         }
 
-        self._unload(model, refiner, deepcache, scale)
-        self._load_pipeline(kind, model, progress, **pipe_kwargs)
-
-        # error loading model
-        if self.pipe is None:
+        self.unload(model, use_refiner, deepcache_interval, scale)
+
+        Pipeline = Config.PIPELINES[kind]
+        Scheduler = Config.SCHEDULERS[scheduler]
+
+        try:
+            with timer(f"Loading {model}"):
+                self.model = model
+                if model.lower() 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,
+                    ).to("cuda")
+                else:
+                    self.pipeline = Pipeline.from_pretrained(model, **pipe_kwargs).to("cuda")
+        except Exception as e:
+            self.log.error(f"Error loading {model}: {e}")
+            self.model = None
+            self.pipeline = None
             return
 
-        same_scheduler = isinstance(self.pipe.scheduler, Config.SCHEDULERS[scheduler])
+        if not isinstance(self.pipeline, Pipeline):
+            self.pipeline = Pipeline.from_pipe(self.pipeline).to("cuda")
+
+        if self.pipeline is not None:
+            self.pipeline.set_progress_bar_config(disable=progress is not None)
+
+        # Check and update scheduler if necessary
+        same_scheduler = isinstance(self.pipeline.scheduler, Scheduler)
         same_karras = (
-            not hasattr(self.pipe.scheduler.config, "use_karras_sigmas")
-            or self.pipe.scheduler.config.use_karras_sigmas == karras
+            not hasattr(self.pipeline.scheduler.config, "use_karras_sigmas")
+            or self.pipeline.scheduler.config.use_karras_sigmas == use_karras
         )
 
-        # same model, different scheduler
         if self.model.lower() == model.lower():
             if not same_scheduler:
                 self.log.info(f"Enabling {scheduler}")
             if not same_karras:
-                self.log.info(f"{'Enabling' if karras else 'Disabling'} Karras sigmas")
+                self.log.info(f"{'Enabling' if use_karras else 'Disabling'} Karras sigmas")
             if not same_scheduler or not same_karras:
-                self.pipe.scheduler = Config.SCHEDULERS[scheduler](**scheduler_kwargs)
+                self.pipeline.scheduler = Scheduler(**scheduler_kwargs)
                 if self.refiner is not None:
-                    self.refiner.scheduler = self.pipe.scheduler
+                    self.refiner.scheduler = self.pipeline.scheduler
 
-        if self._should_load_refiner(refiner):
-            # https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/model_index.json
+        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.pipe.vae,
-                "scheduler": self.pipe.scheduler,
-                "tokenizer_2": self.pipe.tokenizer_2,
-                "text_encoder_2": self.pipe.text_encoder_2,
+                "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, progress, **refiner_kwargs)  # load refiner before deepcache
+            self.load_refiner(refiner_kwargs, progress)
 
-        if self._should_load_deepcache(deepcache):
-            self._load_deepcache(deepcache)
+        if self.should_load_deepcache(deepcache_interval):
+            self.load_deepcache(deepcache_interval)
 
-        if self._should_load_upscaler(scale):
-            self._load_upscaler(scale)
+        if self.should_load_upscaler(scale):
+            self.load_upscaler(scale)

lib/utils.py

@@ -1,24 +1,13 @@
 import functools
-import inspect
 import json
 import time
 from contextlib import contextmanager
-from typing import Callable, TypeVar
 
-import anyio
 import torch
-from anyio import Semaphore
 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
-from typing_extensions import ParamSpec
-
-T = TypeVar("T")
-P = ParamSpec("P")
-
-MAX_CONCURRENT_THREADS = 1
-MAX_THREADS_GUARD = Semaphore(MAX_CONCURRENT_THREADS)
 
 
 @contextmanager
@@ -61,7 +50,7 @@ def safe_progress(progress, current=0, total=0, desc=""):
         progress((current, total), desc=desc)
 
 
-def clear_cuda_cache():
+def cuda_collect():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
@@ -83,14 +72,3 @@ def download_repo_files(repo_id, allow_patterns, token=None):
     if was_disabled:
         disable_progress_bars()
     return snapshot_path
-
-
-# Like the original but supports args and kwargs instead of a dict
-# https://github.com/huggingface/huggingface-inference-toolkit/blob/0.2.0/src/huggingface_inference_toolkit/async_utils.py
-async def async_call(fn: Callable[P, T], *args: P.args, **kwargs: P.kwargs) -> T:
-    async with MAX_THREADS_GUARD:
-        sig = inspect.signature(fn)
-        bound_args = sig.bind(*args, **kwargs)
-        bound_args.apply_defaults()
-        partial_fn = functools.partial(fn, **bound_args.arguments)
-        return await anyio.to_thread.run_sync(partial_fn)

requirements.txt

@@ -1,4 +1,3 @@
-anyio==4.6.1
 compel==2.0.3
 deepcache==0.1.1
 diffusers==0.30.3
@@ -8,5 +7,5 @@ hf-transfer
 numpy==1.26.4
 ruff==0.6.9
 spaces==0.30.4
-torch==2.2.0
-torchvision==0.17.0
+torch==2.4.0
+torchvision==0.19.0