Added model revision as a settable environment variable for snapshot_download

Author: Yushi Homma

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")

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':