Apply style fixes

Author: github-actions[bot]

examples/server-async/utils/requestscopedpipeline.py

@@ -1,17 +1,22 @@
-from typing import Optional, Any, Iterable, List
 import copy
 import threading
+from typing import Any, Iterable, List, Optional
+
 import torch
+
 from diffusers.utils import logging
+
 from .scheduler import BaseAsyncScheduler, async_retrieve_timesteps
-from .wrappers import ThreadSafeTokenizerWrapper, ThreadSafeVAEWrapper, ThreadSafeImageProcessorWrapper
+from .wrappers import ThreadSafeImageProcessorWrapper, ThreadSafeTokenizerWrapper, ThreadSafeVAEWrapper
+
 
 logger = logging.get_logger(__name__)
 
+
 class RequestScopedPipeline:
     DEFAULT_MUTABLE_ATTRS = [
         "_all_hooks",
-        "_offload_device", 
+        "_offload_device",
         "_progress_bar_config",
         "_progress_bar",
         "_rng_state",
@@ -29,42 +34,39 @@ def __init__(
         wrap_scheduler: bool = True,
     ):
         self._base = pipeline
-        
-        
+
         self.unet = getattr(pipeline, "unet", None)
-        self.vae = getattr(pipeline, "vae", None) 
+        self.vae = getattr(pipeline, "vae", None)
         self.text_encoder = getattr(pipeline, "text_encoder", None)
         self.components = getattr(pipeline, "components", None)
-        
+
         self.transformer = getattr(pipeline, "transformer", None)
-        
-        if wrap_scheduler and hasattr(pipeline, 'scheduler') and pipeline.scheduler is not 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()
 
         self._vae_lock = threading.Lock()
         self._image_lock = threading.Lock()
-        
+
         self._auto_detect_mutables = bool(auto_detect_mutables)
         self._tensor_numel_threshold = int(tensor_numel_threshold)
         self._auto_detected_attrs: List[str] = []
 
     def _detect_kernel_pipeline(self, pipeline) -> bool:
         kernel_indicators = [
-            'text_encoding_cache',
-            'memory_manager', 
-            'enable_optimizations',
-            '_create_request_context',
-            'get_optimization_stats'
+            "text_encoding_cache",
+            "memory_manager",
+            "enable_optimizations",
+            "_create_request_context",
+            "get_optimization_stats",
         ]
-        
-        return any(hasattr(pipeline, attr) for attr in kernel_indicators)
 
+        return any(hasattr(pipeline, attr) for attr in kernel_indicators)
 
     def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str] = None, **clone_kwargs):
         base_sched = getattr(self._base, "scheduler", None)
@@ -78,14 +80,12 @@ def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str]
 
         try:
             return wrapped_scheduler.clone_for_request(
-                num_inference_steps=num_inference_steps, 
-                device=device, 
-                **clone_kwargs
+                num_inference_steps=num_inference_steps, device=device, **clone_kwargs
             )
         except Exception as e:
             logger.debug(f"clone_for_request failed: {e}; trying shallow copy fallback")
             try:
-                if hasattr(wrapped_scheduler, 'scheduler'):
+                if hasattr(wrapped_scheduler, "scheduler"):
                     try:
                         copied_scheduler = copy.copy(wrapped_scheduler.scheduler)
                         return BaseAsyncScheduler(copied_scheduler)
@@ -95,8 +95,10 @@ def _make_local_scheduler(self, num_inference_steps: int, device: Optional[str]
                     copied_scheduler = copy.copy(wrapped_scheduler)
                     return BaseAsyncScheduler(copied_scheduler)
             except Exception as e2:
-                logger.warning(f"Shallow copy of scheduler also failed: {e2}. Using original scheduler (*thread-unsafe but functional*).")
-                return wrapped_scheduler 
+                logger.warning(
+                    f"Shallow copy of scheduler also failed: {e2}. Using original scheduler (*thread-unsafe but functional*)."
+                )
+                return wrapped_scheduler
 
     def _autodetect_mutables(self, max_attrs: int = 40):
         if not self._auto_detect_mutables:
@@ -107,16 +109,15 @@ def _autodetect_mutables(self, max_attrs: int = 40):
 
         candidates: List[str] = []
         seen = set()
-        
-        
+
         for name in dir(self._base):
             if name.startswith("__"):
                 continue
             if name in self._mutable_attrs:
                 continue
             if name in ("to", "save_pretrained", "from_pretrained"):
                 continue
-                
+
             try:
                 val = getattr(self._base, name)
             except Exception:
@@ -165,7 +166,9 @@ def _clone_mutable_attrs(self, base, local):
         attrs_to_clone = list(self._mutable_attrs)
         attrs_to_clone.extend(self._autodetect_mutables())
 
-        EXCLUDE_ATTRS = {"components",}
+        EXCLUDE_ATTRS = {
+            "components",
+        }
 
         for attr in attrs_to_clone:
             if attr in EXCLUDE_ATTRS:
@@ -213,16 +216,16 @@ def _clone_mutable_attrs(self, base, local):
     def _is_tokenizer_component(self, component) -> bool:
         if component is None:
             return False
-        
-        tokenizer_methods = ['encode', 'decode', 'tokenize', '__call__']
+
+        tokenizer_methods = ["encode", "decode", "tokenize", "__call__"]
         has_tokenizer_methods = any(hasattr(component, method) for method in tokenizer_methods)
-        
+
         class_name = component.__class__.__name__.lower()
-        has_tokenizer_in_name = 'tokenizer' in class_name
-        
-        tokenizer_attrs = ['vocab_size', 'pad_token', 'eos_token', 'bos_token']
+        has_tokenizer_in_name = "tokenizer" in class_name
+
+        tokenizer_attrs = ["vocab_size", "pad_token", "eos_token", "bos_token"]
         has_tokenizer_attrs = any(hasattr(component, attr) for attr in tokenizer_attrs)
-        
+
         return has_tokenizer_methods and (has_tokenizer_in_name or has_tokenizer_attrs)
 
     def _should_wrap_tokenizers(self) -> bool:
@@ -238,11 +241,21 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] =
             local_pipe = copy.deepcopy(self._base)
 
         try:
-            if hasattr(local_pipe, "vae") and local_pipe.vae is not None and not isinstance(local_pipe.vae, ThreadSafeVAEWrapper):
+            if (
+                hasattr(local_pipe, "vae")
+                and local_pipe.vae is not None
+                and not isinstance(local_pipe.vae, ThreadSafeVAEWrapper)
+            ):
                 local_pipe.vae = ThreadSafeVAEWrapper(local_pipe.vae, self._vae_lock)
 
-            if hasattr(local_pipe, "image_processor") and local_pipe.image_processor is not None and not isinstance(local_pipe.image_processor, ThreadSafeImageProcessorWrapper):
-                local_pipe.image_processor = ThreadSafeImageProcessorWrapper(local_pipe.image_processor, self._image_lock)
+            if (
+                hasattr(local_pipe, "image_processor")
+                and local_pipe.image_processor is not None
+                and not isinstance(local_pipe.image_processor, ThreadSafeImageProcessorWrapper)
+            ):
+                local_pipe.image_processor = ThreadSafeImageProcessorWrapper(
+                    local_pipe.image_processor, self._image_lock
+                )
         except Exception as e:
             logger.debug(f"Could not wrap vae/image_processor: {e}")
 
@@ -253,7 +266,7 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] =
                     num_inference_steps=num_inference_steps,
                     device=device,
                     return_scheduler=True,
-                    **{k: v for k, v in kwargs.items() if k in ['timesteps', 'sigmas']}
+                    **{k: v for k, v in kwargs.items() if k in ["timesteps", "sigmas"]},
                 )
 
                 final_scheduler = BaseAsyncScheduler(configured_scheduler)
@@ -262,10 +275,9 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] =
                 logger.warning("Could not set scheduler on local pipe; proceeding without replacing scheduler.")
 
         self._clone_mutable_attrs(self._base, local_pipe)
-        
 
         original_tokenizers = {}
-        
+
         if self._should_wrap_tokenizers():
             try:
                 for name in dir(local_pipe):
@@ -281,7 +293,7 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] =
                     for key, val in local_pipe.components.items():
                         if val is None:
                             continue
-                        
+
                         if self._is_tokenizer_component(val):
                             if not isinstance(val, ThreadSafeTokenizerWrapper):
                                 original_tokenizers[f"components[{key}]"] = val
@@ -293,9 +305,8 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] =
 
         result = None
         cm = getattr(local_pipe, "model_cpu_offload_context", None)
-        
+
         try:
-            
             if callable(cm):
                 try:
                     with cm():
@@ -316,8 +327,8 @@ def generate(self, *args, num_inference_steps: int = 50, device: Optional[str] =
             try:
                 for name, tok in original_tokenizers.items():
                     if name.startswith("components["):
-                        key = name[len("components["):-1]
-                        if hasattr(local_pipe, 'components') and isinstance(local_pipe.components, dict):
+                        key = name[len("components[") : -1]
+                        if hasattr(local_pipe, "components") and isinstance(local_pipe.components, dict):
                             local_pipe.components[key] = tok
                     else:
                         setattr(local_pipe, name, tok)

examples/server-async/utils/wrappers.py

@@ -4,31 +4,38 @@ def __init__(self, tokenizer, lock):
         self._lock = lock
 
         self._thread_safe_methods = {
-            '__call__', 'encode', 'decode', 'tokenize', 
-            'encode_plus', 'batch_encode_plus', 'batch_decode'
+            "__call__",
+            "encode",
+            "decode",
+            "tokenize",
+            "encode_plus",
+            "batch_encode_plus",
+            "batch_decode",
         }
-    
+
     def __getattr__(self, name):
         attr = getattr(self._tokenizer, name)
-        
+
         if name in self._thread_safe_methods and callable(attr):
+
             def wrapped_method(*args, **kwargs):
                 with self._lock:
                     return attr(*args, **kwargs)
+
             return wrapped_method
-        
+
         return attr
 
     def __call__(self, *args, **kwargs):
         with self._lock:
             return self._tokenizer(*args, **kwargs)
-    
+
     def __setattr__(self, name, value):
-        if name.startswith('_'):
+        if name.startswith("_"):
             super().__setattr__(name, value)
         else:
             setattr(self._tokenizer, name, value)
-    
+
     def __dir__(self):
         return dir(self._tokenizer)
 
@@ -41,9 +48,11 @@ def __init__(self, vae, lock):
     def __getattr__(self, name):
         attr = getattr(self._vae, name)
         if name in {"decode", "encode", "forward"} and callable(attr):
+
             def wrapped(*args, **kwargs):
                 with self._lock:
                     return attr(*args, **kwargs)
+
             return wrapped
         return attr
 
@@ -53,6 +62,7 @@ def __setattr__(self, name, value):
         else:
             setattr(self._vae, name, value)
 
+
 class ThreadSafeImageProcessorWrapper:
     def __init__(self, proc, lock):
         self._proc = proc
@@ -61,14 +71,16 @@ def __init__(self, proc, lock):
     def __getattr__(self, name):
         attr = getattr(self._proc, name)
         if name in {"postprocess", "preprocess"} and callable(attr):
+
             def wrapped(*args, **kwargs):
                 with self._lock:
                     return attr(*args, **kwargs)
+
             return wrapped
         return attr
 
     def __setattr__(self, name, value):
         if name.startswith("_"):
             super().__setattr__(name, value)
         else:
-            setattr(self._proc, name, value)
+            setattr(self._proc, name, value)