Added model revision as a settable environment variable for snapshot_…

…download
hommayushi3 · Jul 25, 2023 · b765dce · b765dce
1 parent e24c26a
commit b765dce
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/handler.py b/handler.py
@@ -32,7 +32,7 @@ def load_model():
     global generator, default_settings
 
     if not generator:
-        model_directory = snapshot_download(repo_id=os.environ["MODEL_REPO"])
+        model_directory = snapshot_download(repo_id=os.environ["MODEL_REPO"], revision=os.getenv("MODEL_REVISION", "main"))
         tokenizer_path = os.path.join(model_directory, "tokenizer.model")
         model_config_path = os.path.join(model_directory, "config.json")
         st_pattern = os.path.join(model_directory, "*.safetensors")

diff --git a/predict.py b/predict.py
@@ -54,7 +54,6 @@ def stream_output(task_id, stream=False):
                         sys.stdout.write(new_output[len(previous_output):])
                         sys.stdout.flush()
                         previous_output = new_output
-                elif len(data['stream']) > 0:
                     return data['stream'][0]['output']
 
                 if data.get('status') == 'COMPLETED':