Update app.py

app.py

@@ -1,9 +1,8 @@
 import gradio as gr
 import torch
-from PIL import Image
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_NAME = "openbmb/MiniCPM-V-2_6-int4"
+MODEL_NAME = "openbmb/MiniCPM-2B-sft-bf16"
 
 print("Loading model...")
 
@@ -12,7 +11,7 @@ tokenizer = AutoTokenizer.from_pretrained(
     trust_remote_code=True
 )
 
-model = AutoModel.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     trust_remote_code=True,
     torch_dtype=torch.float32,
@@ -21,39 +20,46 @@ model = AutoModel.from_pretrained(
 
 model.eval()
 
-def chat_with_image(image, prompt):
+def generate_response(prompt, max_new_tokens=256):
 
-    if image is None:
-        return "Upload an image first."
-
-    image = Image.fromarray(image)
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt"
+    )
 
-    msgs = [
-        {
-            "role": "user",
-            "content": [image, prompt]
-        }
-    ]
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )
 
-    response = model.chat(
-        image=None,
-        msgs=msgs,
-        tokenizer=tokenizer
+    response = tokenizer.decode(
+        outputs[0],
+        skip_special_tokens=True
     )
 
     return response
 
 demo = gr.Interface(
-    fn=chat_with_image,
+    fn=generate_response,
     inputs=[
-        gr.Image(type="numpy", label="Image"),
         gr.Textbox(
-            value="Describe this image in detail.",
-            label="Prompt"
+            lines=6,
+            placeholder="Type your prompt here..."
+        ),
+        gr.Slider(
+            minimum=32,
+            maximum=1024,
+            value=256,
+            step=32,
+            label="Max New Tokens"
         )
     ],
-    outputs=gr.Textbox(label="Response"),
-    title="MiniCPM-V 2.6 CPU"
+    outputs=gr.Textbox(lines=12),
+    title="MiniCPM Text Generation (CPU)"
 )
 
 demo.launch(