From e57d6ebe149fd50daf352df527ca700f00425bb1 Mon Sep 17 00:00:00 2001
From: wangzhihong <wangzhihong@sensetime.com>
Date: Wed, 12 Jun 2024 18:10:58 +0800
Subject: [PATCH] fix pipeline.bind bug

---
 lazyllm/common/common.py                         |  6 ------
 lazyllm/components/auto/autodeploy.py            |  1 -
 .../utils/downloader/model_downloader.py         | 10 +++++-----
 lazyllm/flow.py                                  | 16 ++++++++++++++--
 4 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/lazyllm/common/common.py b/lazyllm/common/common.py
index 0dd1e6cb..7e4cbb66 100644
--- a/lazyllm/common/common.py
+++ b/lazyllm/common/common.py
@@ -183,12 +183,6 @@ def __setattr__(self, __name: str, __value: Any) -> None:
             return setattr(self._f, __name, __value)
         return super(__class__, self).__setattr__(__name, __value)
 
-    def __enter__(self):
-        self._f.__enter__()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        return self._f.__exit__(exc_type, exc_val, exc_tb)
 
 setattr(builtins, 'bind', Bind)
 
diff --git a/lazyllm/components/auto/autodeploy.py b/lazyllm/components/auto/autodeploy.py
index 29bccbf2..432afb89 100644
--- a/lazyllm/components/auto/autodeploy.py
+++ b/lazyllm/components/auto/autodeploy.py
@@ -30,4 +30,3 @@ def __new__(cls, base_model, source=lazyllm.config['model_source'], trust_remote
                     kw[value] = getattr(c, key)
             return deploy_cls(trust_remote_code=trust_remote_code, launcher=launcher, stream=stream, **kw)
         raise RuntimeError(f'No valid framework found, candidates are {[c.framework.lower() for c in candidates]}')
-
diff --git a/lazyllm/components/utils/downloader/model_downloader.py b/lazyllm/components/utils/downloader/model_downloader.py
index aae9c7e8..0a7b4413 100644
--- a/lazyllm/components/utils/downloader/model_downloader.py
+++ b/lazyllm/components/utils/downloader/model_downloader.py
@@ -19,14 +19,14 @@ def __init__(self, model_source=lazyllm.config['model_source'],
         self.token = token
         self.cache_dir = cache_dir
         self.model_pathes = model_path.split(":") if len(model_path) > 0 else []
-        
+
     @classmethod
-    def get_model_type(cls, model) ->str:
+    def get_model_type(cls, model) -> str:
         assert isinstance(model, str) and len(model) > 0, "model name should be a non-empty string"
         for name, info in model_name_mapping.items():
             if 'type' not in info: continue
-            
-            model_name_set={name.casefold()}
+
+            model_name_set = {name.casefold()}
             for source in info:
                 if source == 'type': continue
                 model_name_set.add(info[source].split('/')[-1].casefold())
@@ -34,7 +34,7 @@ def get_model_type(cls, model) ->str:
             if model.split(os.sep)[-1].casefold() in model_name_set:
                 return info['type']
         return 'llm'
-            
+
     def download(self, model=''):
         assert isinstance(model, str), "model name should be a string."
         if len(model) == 0 or model[0] in (os.sep, '.', '~'): return model  # Dummy or local model.
diff --git a/lazyllm/flow.py b/lazyllm/flow.py
index fae83f58..8d2ca518 100644
--- a/lazyllm/flow.py
+++ b/lazyllm/flow.py
@@ -36,9 +36,9 @@ def _add(self, k, v):
             else:
                 lazyllm.LOG.warning(f'{k} is already defined in this scope, ignor it')
 
-    def __enter__(self):
+    def __enter__(self, __frame=None):
         assert len(self._items) == 0, f'Cannot init {self.__class__} with items if you want to use it by context.'
-        self._curr_frame = inspect.currentframe().f_back
+        self._curr_frame = __frame if __frame else inspect.currentframe().f_back
         if self._auto_capture:
             self._frame_keys = list(self._curr_frame.f_locals.keys())
         else:
@@ -85,6 +85,18 @@ def for_each(self, filter, action):
                 action(item)
 
 
+def _bind_enter(self):
+    assert isinstance(self._f, FlowBase)
+    self._f.__enter__(inspect.currentframe().f_back)
+    return self
+
+def _bind_exit(self, exc_type, exc_val, exc_tb):
+    return self._f.__exit__(exc_type, exc_val, exc_tb)
+
+setattr(bind, '__enter__', _bind_enter)
+setattr(bind, '__exit__', _bind_exit)
+
+
 def _is_function(f):
     return isinstance(f, (types.BuiltinFunctionType, types.FunctionType,
                           types.BuiltinMethodType, types.MethodType, types.LambdaType))