Update app.py

app.py

@@ -1,112 +1,81 @@
-import spaces  # 必须在最顶部导入
 import gradio as gr
 import os
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, LlavaNextProcessor, LlavaNextForConditionalGeneration
+from PIL import Image
 
 # 获取 Hugging Face 访问令牌
 hf_token = os.getenv("HF_API_TOKEN")
 
-# 定义基础模型名称
-base_model_name = "unsloth/meta-llama-3.1-8b-bnb-4bit"
-
-# 定义 adapter 模型名称
-adapter_model_name = "larry1129/WooWoof_AI"
-
-# 定义全局变量用于缓存模型和分词器
-model = None
-tokenizer = None
+# 定义模型名称
+vqa_model_name = "llava-hf/llava-v1.6-mistral-7b-hf"
+language_model_name = "larry1129/WooWoof_AI_Vision_merged_16bit_3b"
+
+# 全局变量用于缓存模型和分词器
+vqa_processor = None
+vqa_model = None
+language_tokenizer = None
+language_model = None
+
+# 初始化看图说话模型
+def load_vqa_model():
+    global vqa_processor, vqa_model
+    if vqa_processor is None or vqa_model is None:
+        vqa_processor = LlavaNextProcessor.from_pretrained(vqa_model_name, use_auth_token=hf_token)
+        vqa_model = LlavaNextForConditionalGeneration.from_pretrained(
+            vqa_model_name,
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True
+        ).to("cuda:0")
+    return vqa_processor, vqa_model
+
+# 初始化纯语言模型
+def load_language_model():
+    global language_tokenizer, language_model
+    if language_tokenizer is None or language_model is None:
+        language_tokenizer = AutoTokenizer.from_pretrained(language_model_name, use_auth_token=hf_token)
+        language_model = AutoModelForCausalLM.from_pretrained(
+            language_model_name,
+            device_map="auto",
+            torch_dtype=torch.float16
+        )
+        language_tokenizer.pad_token = language_tokenizer.eos_token
+        language_model.config.pad_token_id = language_tokenizer.pad_token_id
+        language_model.eval()
+    return language_tokenizer, language_model
+
+# 从图片生成描述
+def generate_image_description(image):
+    vqa_processor, vqa_model = load_vqa_model()
+    conversation = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What is shown in this image?"},
+                {"type": "image"},
+            ],
+        },
+    ]
+    prompt = vqa_processor.apply_chat_template(conversation, add_generation_prompt=True)
+    inputs = vqa_processor(images=image, text=prompt, return_tensors="pt").to("cuda:0")
 
-# 定义提示生成函数
-def generate_prompt(instruction, input_text=""):
-    if input_text:
-        prompt = f"""### Instruction:
+    with torch.no_grad():
+        output = vqa_model.generate(**inputs, max_new_tokens=100)
+    image_description = vqa_processor.decode(output[0], skip_special_tokens=True)
+    return image_description
+
+# 使用纯语言模型生成最终回答
+def generate_language_response(instruction, image_description):
+    language_tokenizer, language_model = load_language_model()
+    prompt = f"""### Instruction:
 {instruction}
 ### Input:
-{input_text}
+{image_description}
 ### Response:
 """
-    else:
-        prompt = f"""### Instruction:
-{instruction}
-### Response:
-"""
-    return prompt
-
-# 定义生成响应的函数，并使用 @spaces.GPU 装饰
-@spaces.GPU(duration=40)  # 建议将 duration 增加到 120
-def generate_response(instruction, input_text):
-    global model, tokenizer
-
-    if model is None:
-        print("开始加载模型...")
-        # 检查 bitsandbytes 是否已安装
-        import importlib.util
-        if importlib.util.find_spec("bitsandbytes") is None:
-            import subprocess
-            subprocess.call(["pip", "install", "--upgrade", "bitsandbytes"])
-
-        try:
-            # 在函数内部导入需要 GPU 的库
-            import torch
-            from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
-            from peft import PeftModel
-
-            # 创建量化配置
-            bnb_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_quant_type="nf4",
-                bnb_4bit_compute_dtype=torch.float16
-            )
-
-            # 加载分词器
-            tokenizer = AutoTokenizer.from_pretrained(base_model_name, use_auth_token=hf_token)
-            print("分词器加载成功。")
-
-            # 加载基础模型
-            base_model = AutoModelForCausalLM.from_pretrained(
-                base_model_name,
-                quantization_config=bnb_config,
-                device_map="auto",
-                use_auth_token=hf_token,
-                trust_remote_code=True
-            )
-            print("基础模型加载成功。")
-
-            # 加载适配器模型
-            model = PeftModel.from_pretrained(
-                base_model,
-                adapter_model_name,
-                torch_dtype=torch.float16,
-                use_auth_token=hf_token
-            )
-            print("适配器模型加载成功。")
-
-            # 设置 pad_token
-            tokenizer.pad_token = tokenizer.eos_token
-            model.config.pad_token_id = tokenizer.pad_token_id
-
-            # 切换到评估模式
-            model.eval()
-            print("模型已切换到评估模式。")
-        except Exception as e:
-            print("加载模型时出错：", e)
-            raise e
-    else:
-        # 在函数内部导入需要的库
-        import torch
-
-    # 检查 model 和 tokenizer 是否已正确加载
-    if model is None or tokenizer is None:
-        print("模型或分词器未正确加载。")
-        raise ValueError("模型或分词器未正确加载。")
-
-    # 生成提示
-    prompt = generate_prompt(instruction, input_text)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
+    inputs = language_tokenizer(prompt, return_tensors="pt").to(language_model.device)
     with torch.no_grad():
-        outputs = model.generate(
+        outputs = language_model.generate(
             input_ids=inputs["input_ids"],
             attention_mask=inputs.get("attention_mask"),
             max_new_tokens=128,
@@ -114,23 +83,28 @@ def generate_response(instruction, input_text):
             top_p=0.95,
             do_sample=True,
         )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = language_tokenizer.decode(outputs[0], skip_special_tokens=True)
     response = response.split("### Response:")[-1].strip()
     return response
 
-# 创建 Gradio 接口
+# 整合的 Gradio 接口函数
+def process_image_and_text(image, instruction):
+    image_description = generate_image_description(image)
+    final_response = generate_language_response(instruction, image_description)
+    return f"图片描述: {image_description}\n\n最终回答: {final_response}"
+
+# 创建 Gradio 界面
 iface = gr.Interface(
-    fn=generate_response,
+    fn=process_image_and_text,
     inputs=[
-        gr.Textbox(lines=2, placeholder="Instruction", label="Instruction"),
-        gr.Textbox(lines=2, placeholder="Input", label="Input (Option)")
+        gr.Image(type="pil", label="上传图片"),
+        gr.Textbox(lines=2, placeholder="Instruction", label="Instruction")
     ],
     outputs="text",
-    title="WooWoof AI",
-    description="Based on LLAMA 3.1 for pet related",
+    title="WooWoof AI - 图片和文本交互",
+    description="输入图片并添加指令，生成基于图片描述的回答。",
     allow_flagging="never"
 )
 
 # 启动 Gradio 接口
-iface.launch(share=True)
-
+iface.launch()