Add support for the Llama-2 model

apps/stable_diffusion/web/ui/stablelm_ui.py

@@ -21,6 +21,7 @@ vicuna_model = 0
 past_key_values = None
 
 model_map = {
+    "llama2": "meta-llama/Llama-2-7b-chat-hf",
     "codegen": "Salesforce/codegen25-7b-multi",
     "vicuna1p3": "lmsys/vicuna-7b-v1.3",
     "vicuna": "TheBloke/vicuna-7B-1.1-HF",
@@ -29,6 +30,15 @@ model_map = {
 
 # NOTE: Each `model_name` should have its own start message
 start_message = {
+    "llama2": (
+        "System: You are a helpful, respectful and honest assistant. Always answer "
+        "as helpfully as possible, while being safe.  Your answers should not "
+        "include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal "
+        "content. Please ensure that your responses are socially unbiased and positive "
+        "in nature. If a question does not make any sense, or is not factually coherent, "
+        "explain why instead of answering something not correct. If you don't know the "
+        "answer to a question, please don't share false information."
+    ),
     "StableLM": (
         "<|SYSTEM|># StableLM Tuned (Alpha version)"
         "\n- StableLM is a helpful and harmless open-source AI language model "
@@ -57,7 +67,7 @@ start_message = {
 def create_prompt(model_name, history):
     system_message = start_message[model_name]
 
-    if model_name in ["StableLM", "vicuna", "vicuna1p3"]:
+    if model_name in ["StableLM", "vicuna", "vicuna1p3", "llama2"]:
         conversation = "".join(
             [
                 "".join(["<|USER|>" + item[0], "<|ASSISTANT|>" + item[1]])
@@ -86,7 +96,7 @@ def chat(curr_system_message, history, model, device, precision, cli=True):
     global vicuna_model
     model_name, model_path = list(map(str.strip, model.split("=>")))
 
-    if model_name in ["vicuna", "vicuna1p3", "codegen"]:
+    if model_name in ["vicuna", "vicuna1p3", "codegen", "llama2"]:
         from apps.language_models.scripts.vicuna import (
             UnshardedVicuna,
         )