use transform

app.py

@@ -1,64 +1,54 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from PIL import Image
 import base64
+from io import BytesIO
 
-client = InferenceClient("openbmb/MiniCPM-Llama3-V-2_5-int4",trust_remote_code=True)
+# 加载模型和分词器
+model_name = "openbmb/MiniCPM-Llama3-V-2_5-int4"
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
 
-def encode_image(image_path):
-    with open(image_path, "rb") as image_file:
-        return base64.b64encode(image_file.read()).decode('utf-8')
+def encode_image(image):
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-def respond(
-    message,
-    image,
-    history,
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": bot_msg})
-    
-    if image:
-        base64_image = encode_image(image)
-        image_message = f"<image>{base64_image}</image>"
-        message = image_message + "\n" + message
+def generate_text(prompt, max_length=100):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=max_length, num_return_sequences=1)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    messages.append({"role": "user", "content": message})
-    
-    response = ""
-    for message in client.text_generation(
-        prompt=f"{messages}",
-        max_new_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.token.text
-        response += token
-        yield response, history + [(message, response)]
+def predict(image, prompt):
+    if image is not None:
+        # 确保image是PIL Image对象
+        if isinstance(image, str):
+            image = Image.open(image)
+        
+        # 编码图像
+        encoded_image = encode_image(image)
+        
+        # 准备输入
+        full_prompt = f"<image>{encoded_image}</image>\n{prompt if prompt else 'Describe this image.'}"
+        
+        # 生成文本
+        result = generate_text(full_prompt)
+        
+        return f"Model response: {result}\n\nUser prompt: {prompt}"
+    else:
+        return "No image uploaded. " + (f"You asked: {prompt}" if prompt else "Please upload an image and optionally provide a prompt.")
 
 demo = gr.Interface(
-    respond,
+    predict,
     inputs=[
-        gr.Textbox(label="Message"),
-        gr.Image(type="filepath", label="Upload Image"),
-        gr.State([]),  # for history
-        gr.Textbox(value="You are a friendly AI assistant capable of understanding images and text.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Textbox(label="Prompt (optional)")
     ],
-    outputs=[
-        gr.Textbox(label="Response"),
-        gr.State()  # for updated history
-    ],
-    title="MiniCPM-Llama3-V-2_5 Image and Text Chat",
-    description="Upload an image and ask questions about it, or just chat without an image.",
-    allow_flagging="never"
+    outputs=gr.Textbox(label="Result"),
+    title="Image Analysis with MiniCPM-Llama3-V-2_5-int4",
+    description="Upload an image and optionally provide a prompt for analysis."
 )
 
 if __name__ == "__main__":