Update app.py

app.py

@@ -1,57 +1,103 @@
-# 导入库
-from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
+import gradio as gr
+from transformers import pipeline
 from PIL import Image
 import torch
-from transformers import pipeline
-import requests
-from io import BytesIO
 
-# 1. 图像标题生成（使用指定模型）
-def generate_caption(image_url):
-    model_name = "bipin/image-caption-generator"
-    model = VisionEncoderDecoderModel.from_pretrained(model_name)
-    feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
-    tokenizer = AutoTokenizer.from_pretrained("gpt2")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    
-    # 下载并预处理图像
-    response = requests.get(image_url)
-    img = Image.open(BytesIO(response.content)).convert("RGB")
-    pixel_values = feature_extractor(images=[img], return_tensors="pt").pixel_values.to(device)
-    
-    # 生成标题（限制50字内）
-    output_ids = model.generate(pixel_values, num_beams=4, max_length=50)
-    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return caption
+# 初始化模型（缓存加载）
+def init_models():
+    # 鸟类分类模型
+    classifier = pipeline(
+        "image-classification",
+        model="chriamue/bird-species-classifier",
+        device=0 if torch.cuda.is_available() else -1
+    )
+    
+    # 文本生成模型（设置量化降低显存占用）
+    text_generator = pipeline(
+        "text-generation",
+        model="Qwen/Qwen3-235B-A22B",
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+        model_kwargs={"load_in_4bit": True}
+    )
+    
+    # 语音合成模型
+    tts = pipeline(
+        "text-to-speech",
+        model="facebook/mms-tts-eng",
+        device=0 if torch.cuda.is_available() else -1
+    )
+    
+    return classifier, text_generator, tts
+
+# 生成儿童友好的鸟类描述
+def generate_child_friendly_text(bird_name):
+    PROMPT = f"""请用简单易懂的语言，向6-12岁儿童介绍{bird_name}：
+1. 用比喻手法描述外形特征
+2. 解释生活习性时使用拟人化
+3. 包含一个有趣的小知识
+4. 语句长度不超过15个英文单词
+5. 避免使用专业术语"""
+    
+    response = text_generator(
+        PROMPT,
+        max_new_tokens=200,
+        temperature=0.7,
+        do_sample=True
+    )
+    
+    return response[0]['generated_text'].split('\n')[2:]  # 提取核心内容
 
-# 2. 标题扩写为宣传文案（使用文本生成模型）
-def expand_to_copy(caption):
-    generator = pipeline("text-generation", model="gpt2", max_length=200)
-    prompt = f"根据以下图片标题生成宣传文案：{caption}\n要求：生动形象，突出产品优势，适合社交媒体传播。"
-    copy = generator(prompt, num_return_sequences=1)[0]['generated_text']
-    return copy.strip()
+# 主处理流程
+def process_image(image):
+    try:
+        # Step 1: 鸟类识别
+        classification = classifier(image)
+        bird_name = classification[0]['label']
+        
+        # Step 2: 生成描述
+        description = generate_child_friendly_text(bird_name)
+        
+        # Step 3: 语音合成
+        speech = tts(description, forward_params={"speaker_id": 6})  # 使用儿童语音
+        
+        return {
+            "bird_name": bird_name,
+            "description": "\n".join(description),
+            "audio": speech["audio"]
+        }
+    except Exception as e:
+        return f"处理错误: {str(e)}"
 
-# 3. 文本转语音（使用TTS模型）
-def text_to_speech(text, output_file="output.mp3"):
-    tts = pipeline("text-to-speech", model="facebook/t5-small")
-    speech = tts(text)
-    with open(output_file, "wb") as f:
-        f.write(speech["audio"])
-    return output_file
+# 初始化模型
+classifier, text_generator, tts = init_models()
 
-# 主函数
-def marketing_pipeline(image_url):
-    # 生成标题
-    caption = generate_caption(image_url)
-    print(f"生成标题：{caption}")
+# 创建Gradio界面
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🐦 鸟类知识小课堂")
     
-    # 扩写文案
-    copy = expand_to_copy(caption)
-    print(f"宣传文案：\n{copy}")
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="上传鸟类图片")
+        audio_output = gr.Audio(label="语音讲解", autoplay=True)
     
-    # 生成语音
-    audio_file = text_to_speech(copy)
-    print(f"语音文件已保存：{audio_file}")
+    with gr.Column():
+        name_output = gr.Textbox(label="识别到的鸟类")
+        text_output = gr.Textbox(label="趣味知识", lines=4)
     
-    return caption, copy, audio_file
+    examples = gr.Examples(
+        examples=["eagle.jpg", "penguin.jpg", "peacock.jpg"],
+        inputs=image_input
+    )
+    
+    image_input.change(
+        process_image,
+        inputs=image_input,
+        outputs=[name_output, text_output, audio_output]
+    )
+
+# 部署配置
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=True
+)