Update app.py

app.py

@@ -21,6 +21,12 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ctx_limit = 3500
 ########################## text rwkv ################################################################
 from rwkv.utils import PIPELINE, PIPELINE_ARGS
+
+title_v6 = "RWKV-x060-World-1B6-v2-20240208-ctx4096"
+model_path_v6 = hf_hub_download(repo_id="BlinkDL/rwkv-6-world", filename=f"{title_v6}.pth")
+model_v6 = RWKV(model=model_path_v6, strategy='cuda fp16')
+pipeline_v6 = PIPELINE(model_v6, "rwkv_vocab_v20230424")
+
 title = "RWKV-5-World-1B5-v2-20231025-ctx4096"
 model_path = hf_hub_download(repo_id="BlinkDL/rwkv-5-world", filename=f"{title}.pth")
 model = RWKV(model=model_path, strategy='cuda fp16')
@@ -65,7 +71,7 @@ def evaluate(
     state = None
     for i in range(int(token_count)):
         input_ids = pipeline.encode(ctx)[-ctx_limit:] if i == 0 else [token]
-        out, state = model.forward(tokens=input_ids, state=state)
+        out, state = model_v6.forward(tokens=input_ids, state=state)
         for n in occurrence:
             out[n] -= (args.alpha_presence + occurrence[n] * args.alpha_frequency)
 
@@ -261,9 +267,9 @@ def chatbot(image, question):
 
 ##################################################################################################################
 with gr.Blocks(title=title) as demo:
-    gr.HTML(f"<div style=\"text-align: center;\">\n<h1>RWKV-5 World v2 - {title}</h1>\n</div>")
+    gr.HTML(f"<div style=\"text-align: center;\">\n<h1>{title_v6}</h1>\n</div>")
     with gr.Tab("Raw Generation"):
-        gr.Markdown(f"This is [RWKV-5 World v2](https://huggingface.co/BlinkDL/rwkv-5-world) with 1.5B params - a 100% attention-free RNN [RWKV-LM](https://github.com/BlinkDL/RWKV-LM). Supports all 100+ world languages and code. And we have [200+ Github RWKV projects](https://github.com/search?o=desc&p=1&q=rwkv&s=updated&type=Repositories). *** Please try examples first (bottom of page) *** (edit them to use your question). Demo limited to ctxlen {ctx_limit}.")
+        gr.Markdown(f"This is [RWKV-6 World v2](https://huggingface.co/BlinkDL/rwkv-6-world) with 1.6B params - a 100% attention-free RNN [RWKV-LM](https://github.com/BlinkDL/RWKV-LM). Supports all 100+ world languages and code. And we have [200+ Github RWKV projects](https://github.com/search?o=desc&p=1&q=rwkv&s=updated&type=Repositories). *** Please try examples first (bottom of page) *** (edit them to use your question). Demo limited to ctxlen {ctx_limit}.")
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(lines=2, label="Prompt", value="Assistant: Sure! Here is a very detailed plan to create flying pigs:")