Update app.py

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 import torch
+from functools import lru_cache
 
 MODEL_NAME = "openbmb/MiniCPM-V-4"
 
-@gr.cache(allow_output_mutation=True)
-def load_model():
+@lru_cache(maxsize=1)
+def load_pipeline():
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
@@ -13,15 +14,14 @@ def load_model():
         device_map="auto",
         trust_remote_code=True
     )
-    pipeline = TextGenerationPipeline(
+    return TextGenerationPipeline(
         model=model,
         tokenizer=tokenizer,
         device=model.device.index if torch.cuda.is_available() else -1
     )
-    return pipeline
 
-def generate(prompt: str, max_length: int = 100, top_k: int = 50, top_p: float = 0.95):
-    pipe = load_model()
+def generate(prompt: str, max_length: int, top_k: int, top_p: float):
+    pipe = load_pipeline()
     output = pipe(
         prompt,
         max_length=max_length,
@@ -35,14 +35,10 @@ def generate(prompt: str, max_length: int = 100, top_k: int = 50, top_p: float =
 with gr.Blocks() as demo:
     gr.Markdown("# MiniCPM-V-4 Text Generation Demo")
 
-    with gr.Row():
-        prompt_input = gr.Textbox(label="Prompt", placeholder="여기에 입력하세요...", lines=2)
-
-    with gr.Row():
-        max_len = gr.Slider(10, 512, value=100, step=10, label="Max Length")
-        topk = gr.Slider(1, 100, value=50, step=1, label="Top-k")
-        topp = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
-
+    prompt_input = gr.Textbox(label="Prompt", lines=2, placeholder="여기에 입력하세요...")
+    max_len = gr.Slider(10, 512, value=100, step=10, label="Max Length")
+    topk = gr.Slider(1, 100, value=50, step=1, label="Top-k")
+    topp = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
     generate_btn = gr.Button("Generate")
     output_box = gr.Textbox(label="Generated Text", lines=5)