Update app.py

app.py

@@ -31,6 +31,16 @@ pipeline_v6 = PIPELINE(model_v6, "rwkv_vocab_v20230424")
 
 args = model_v6.args
 
+_, _ = model.forward([0], None)
+state = model.generate_zero_state()
+static_input = torch.empty((model.n_embd), device="cuda", dtype=torch.half)
+static_state_in = [torch.empty_like(x, device="cuda") for x in state]
+static_state_out = [torch.empty_like(x, device="cuda") for x in state]
+static_output = torch.empty((model.args.vocab_size), device="cuda", dtype=torch.half)
+graph = torch.cuda.CUDAGraph()
+with torch.cuda.graph(graph):
+    static_output, static_state_out = model.forward_one_alt(static_input, static_state_in)
+
 penalty_decay = 0.996
 
 def generate_prompt(instruction, input=""):