Dynamic Temperature HF loader support (#5174)

---------

Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
This commit is contained in:
kalomaze 2024-01-07 07:36:26 -06:00 committed by GitHub
parent 3eca20c015
commit 48327cc5c4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
14 changed files with 184 additions and 8 deletions

View file

@ -10,6 +10,10 @@ from transformers import is_torch_xpu_available
import modules.shared as shared
class StopNowException(Exception):
pass
class _StopEverythingStoppingCriteria(transformers.StoppingCriteria):
def __init__(self):
transformers.StoppingCriteria.__init__(self)
@ -49,13 +53,13 @@ class Iteratorize:
def _callback(val):
if self.stop_now or shared.stop_everything:
raise ValueError
raise StopNowException
self.q.put(val)
def gentask():
try:
ret = self.mfunc(callback=_callback, *args, **self.kwargs)
except ValueError:
except StopNowException:
pass
except:
traceback.print_exc()