Update examples/server-async/utils/*

Author: FredyRivera-dev

examples/server-async/DiffusersServer/serverasync.py

@@ -5,7 +5,7 @@
 from pydantic import BaseModel
 from .Pipelines import TextToImagePipelineSD3, TextToImagePipelineFlux, TextToImagePipelineSD, logger
 import logging
-from ..utils import RequestScopedPipeline
+from .utils import RequestScopedPipeline
 from diffusers import *
 import random
 import uuid

examples/server-async/utils/__init__.py …-async/DiffusersServer/utils/__init__.pyexamples/server-async/utils/__init__.py renamed to examples/server-async/DiffusersServer/utils/__init__.py examples/server-async/utils/__init__.py renamed to examples/server-async/DiffusersServer/utils/__init__.py

@@ -3,6 +3,8 @@
 import threading
 import torch
 from diffusers.utils import logging
+from .scheduler import BaseAsyncScheduler, async_retrieve_timesteps
+
 
 logger = logging.get_logger(__name__)
 
@@ -27,14 +29,19 @@ def __init__(
         mutable_attrs: Optional[Iterable[str]] = None,
         auto_detect_mutables: bool = True,
         tensor_numel_threshold: int = 1_000_000,
-        tokenizer_lock: Optional[threading.Lock] = None
+        tokenizer_lock: Optional[threading.Lock] = None,
+        wrap_scheduler: bool = True
     ):
         self._base = pipeline
         self.unet = getattr(pipeline, "unet", None)
         self.vae = getattr(pipeline, "vae", None)
         self.text_encoder = getattr(pipeline, "text_encoder", None)
         self.components = getattr(pipeline, "components", None)
 
+        if wrap_scheduler and hasattr(pipeline, 'scheduler') and pipeline.scheduler is not None:
+            if not isinstance(pipeline.scheduler, BaseAsyncScheduler):
+                pipeline.scheduler = BaseAsyncScheduler(pipeline.scheduler)
+
         self._mutable_attrs = list(mutable_attrs) if mutable_attrs is not None else list(self.DEFAULT_MUTABLE_ATTRS)
         self._tokenizer_lock = tokenizer_lock if tokenizer_lock is not None else threading.Lock()
 
@@ -48,17 +55,24 @@ def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str]
         if base_sched is None:
             return None
 
-        if hasattr(base_sched, "clone_for_request"):
-            try:
-                return base_sched.clone_for_request(num_inference_steps=num_inference_steps, device=device, **clone_kwargs)
-            except Exception as e:
-                logger.debug(f"clone_for_request failed: {e}; falling back to deepcopy()")
+        if not isinstance(base_sched, BaseAsyncScheduler):
+            wrapped_scheduler = BaseAsyncScheduler(base_sched)
+        else:
+            wrapped_scheduler = base_sched
 
         try:
-            return copy.deepcopy(base_sched)
+            return wrapped_scheduler.clone_for_request(
+                num_inference_steps=num_inference_steps, 
+                device=device, 
+                **clone_kwargs
+            )
         except Exception as e:
-            logger.warning(f"Deepcopy of scheduler failed: {e}. Returning original scheduler (*risky*).")
-            return base_sched  
+            logger.debug(f"clone_for_request failed: {e}; falling back to deepcopy()")
+            try:
+                return copy.deepcopy(wrapped_scheduler)
+            except Exception as e:
+                logger.warning(f"Deepcopy of scheduler failed: {e}. Returning original scheduler (*risky*).")
+                return wrapped_scheduler  
 
     def _autodetect_mutables(self, max_attrs: int = 40):
         if not self._auto_detect_mutables:
@@ -197,7 +211,16 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] =
 
         if local_scheduler is not None:
             try:
-                setattr(local_pipe, "scheduler", local_scheduler)
+                timesteps, num_steps, configured_scheduler = async_retrieve_timesteps(
+                    local_scheduler.scheduler,
+                    num_inference_steps=num_inference_steps,
+                    device=device,
+                    return_scheduler=True,
+                    **{k: v for k, v in kwargs.items() if k in ['timesteps', 'sigmas']}
+                )
+
+                final_scheduler = BaseAsyncScheduler(configured_scheduler)
+                setattr(local_pipe, "scheduler", final_scheduler)
             except Exception:
                 logger.warning("Could not set scheduler on local pipe; proceeding without replacing scheduler.")
 

…ver-async/utils/requestscopedpipeline.py …ersServer/utils/requestscopedpipeline.pyexamples/server-async/utils/requestscopedpipeline.py renamed to examples/server-async/DiffusersServer/utils/requestscopedpipeline.py examples/server-async/utils/requestscopedpipeline.py renamed to examples/server-async/DiffusersServer/utils/requestscopedpipeline.py

@@ -5,14 +5,16 @@
 
 class BaseAsyncScheduler:
     def __init__(self, scheduler: Any):
-        pass
-
-    def clone_for_request(self, num_inference_steps: int, device: Union[str, torch.device] = None):
-        # I leave it as an example of what the Scheduler should do to implement it later
-        """local = copy.deepcopy(self)
-        local.set_timesteps(num_inference_steps=num_inference_steps, device=device)
-        return local"""
-        pass
+        self.scheduler = scheduler
+
+    def clone_for_request(self, num_inference_steps: int, device: Union[str, torch.device, None] = None, **kwargs):
+        local = copy.deepcopy(self.scheduler)
+
+        local.set_timesteps(num_inference_steps=num_inference_steps, device=device, **kwargs)
+
+        cloned = self.__class__(local)
+        
+        return cloned
 
 
 def async_retrieve_timesteps(