import os
import gradio as gr
from transformers import pipeline, set_seed

# 加载第一个pipeline，用于图片描述
image_to_text_pipe = pipeline("image-to-text",
                              model="Salesforce/blip-image-captioning-base")

# 加载第二个pipeline，用于根据文本生成故事
text_to_story_pipe = pipeline("text-generation",
                              model="mistralai/Mistral-7B-Instruct-v0.2")

# 设置随机种子以保证结果的一致性
set_seed(42)

def generate_description(image):
    # 从图片生成描述
    description = image_to_text_pipe(image)
    description_text = description[0]['generated_text']
    return description_text

def generate_story(text):
    # 根据描述生成故事
    story = text_to_story_pipe(text, max_length=300, num_return_sequences=1)
    story_text = story[0]['generated_text']
    return story_text

def process_image_and_generate_story(image):
    # 串联两个处理过程：生成描述和根据描述生成故事
    description_text = generate_description(image)
    story_text = generate_story(description_text)
    return story_text

# 设置Gradio接口
iface = gr.Interface(fn=process_image_and_generate_story,
                     inputs=gr.Image(type='pil'),
                     outputs="text",
                     title="Image to Story Generator",
                     description="Upload an image and the system will generate a story based on it.")

# 启动界面
iface.launch()