Add wrappers.py

Author: FredyRivera-dev

examples/server-async/utils/requestscopedpipeline.py

@@ -4,85 +4,10 @@
 import torch
 from diffusers.utils import logging
 from .scheduler import BaseAsyncScheduler, async_retrieve_timesteps
+from .wrappers import ThreadSafeTokenizerWrapper, ThreadSafeVAEWrapper, ThreadSafeImageProcessorWrapper
 
 logger = logging.get_logger(__name__)
 
-class ThreadSafeTokenizerWrapper:
-    def __init__(self, tokenizer, lock):
-        self._tokenizer = tokenizer
-        self._lock = lock
-
-        self._thread_safe_methods = {
-            '__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('_'):
-            super().__setattr__(name, value)
-        else:
-            setattr(self._tokenizer, name, value)
-    
-    def __dir__(self):
-        return dir(self._tokenizer)
-
-
-class ThreadSafeVAEWrapper:
-    def __init__(self, vae, lock):
-        self._vae = vae
-        self._lock = lock
-
-    def __getattr__(self, name):
-        attr = getattr(self._vae, name)
-        # métodos que queremos proteger
-        if name in {"decode", "encode", "forward"} 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._vae, name, value)
-
-class ThreadSafeImageProcessorWrapper:
-    def __init__(self, proc, lock):
-        self._proc = proc
-        self._lock = 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)
-
 class RequestScopedPipeline:
     DEFAULT_MUTABLE_ATTRS = [
         "_all_hooks",

examples/server-async/utils/wrappers.py

@@ -0,0 +1,74 @@
+class ThreadSafeTokenizerWrapper:
+    def __init__(self, tokenizer, lock):
+        self._tokenizer = tokenizer
+        self._lock = lock
+
+        self._thread_safe_methods = {
+            '__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('_'):
+            super().__setattr__(name, value)
+        else:
+            setattr(self._tokenizer, name, value)
+    
+    def __dir__(self):
+        return dir(self._tokenizer)
+
+
+class ThreadSafeVAEWrapper:
+    def __init__(self, vae, lock):
+        self._vae = vae
+        self._lock = 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
+
+    def __setattr__(self, name, value):
+        if name.startswith("_"):
+            super().__setattr__(name, value)
+        else:
+            setattr(self._vae, name, value)
+
+class ThreadSafeImageProcessorWrapper:
+    def __init__(self, proc, lock):
+        self._proc = proc
+        self._lock = 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)