From b765dcec3d14c830a8923f7c366a65a4ddbe7551 Mon Sep 17 00:00:00 2001 From: Yushi Homma Date: Mon, 24 Jul 2023 23:46:31 -0700 Subject: [PATCH] Added model revision as a settable environment variable for snapshot_download --- handler.py | 2 +- predict.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/handler.py b/handler.py index 53035e4..4be7828 100644 --- 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 index 93d29ee..880e289 100644 --- 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':