Update app.py

app.py

@@ -1,56 +1,36 @@
-import streamlit as st
-from transformers import pipeline, AutoProcessor, AutoModel
-from scipy.io.wavfile import write as write_wav
-import numpy as np
+import gradio as gr
+from transformers import pipeline
+from diffusers import StableDiffusionPipeline
 import torch
 
-def img2text(image_path):
-    img2caption = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
-    return img2caption(image_path)[0]['generated_text']
-
-def text2story(text):
-    pipe = pipeline("text-generation", model="pranavpsv/genre-story-generator-v2")
-    story_text = pipe(text, max_length=100)[0]['generated_text']
-    return story_text
-
-def text2audio(story_text):
-    processor = AutoProcessor.from_pretrained("facebook/mms-tts-eng")
-    model = AutoModel.from_pretrained("facebook/mms-tts-eng")
-
-    inputs = processor(text=story_text, return_tensors="pt")
-    with torch.no_grad():
-        output = model(**inputs).waveform
-
-    audio_array = output.cpu().numpy().squeeze()
-    sample_rate = 16000 
-
-    return audio_array, sample_rate
-
-# Streamlit UI
-st.set_page_config(page_title="Your Image to Audio Story", page_icon="🦜")
-st.header("Turn Your Image to Audio Story")
-uploaded_file = st.file_uploader("Select an Image...")
-
-if uploaded_file is not None:
-    bytes_data = uploaded_file.getvalue()
-    with open(uploaded_file.name, "wb") as file:
-        file.write(bytes_data)
-    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
-
-    # Stage 1: Image to Text
-    st.text('Processing img2text...')
-    scenario = img2text(uploaded_file.name)
-    st.write(scenario)
-
-    # Stage 2: Text to Story
-    st.text('Generating a story...')
-    story = text2story(scenario)
-    st.write(story)
-
-    # Stage 3: Story to Audio
-    st.text('Generating audio...')
-    audio_array, sample_rate = text2audio(story)
-    
-    audio_file = "output_audio.wav"
-    write_wav(audio_file, sample_rate, audio_array)
-    st.audio(audio_file)
+# 文本生成模型：使用 GPT-2 或其他公开模型
+text_generator = pipeline("text-generation", model="gpt2")
+
+# 图像生成模型：使用 Stable Diffusion
+device = "cuda" if torch.cuda.is_available() else "cpu"
+image_generator = StableDiffusionPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+).to(device)
+
+def generate_text_and_image(keyword):
+    # Step 1: 根据关键词生成文本
+    prompt = f"Write a short paragraph about {keyword}."
+    text_output = text_generator(prompt, max_length=100, do_sample=True, temperature=0.8)[0]["generated_text"]
+
+    # Step 2: 使用生成的文本生成图像
+    image = image_generator(text_output).images[0]
+
+    return text_output, image
+
+# Gradio界面
+iface = gr.Interface(
+    fn=generate_text_and_image,
+    inputs=gr.Textbox(lines=1, placeholder="请输入关键词，例如：a futuristic city"),
+    outputs=[gr.Textbox(label="生成的文字"), gr.Image(label="根据文字生成的图像")],
+    title="关键词到文字 + 图像生成器",
+    description="输入关键词，先生成一段文字，再根据这段文字生成图像（使用Hugging Face模型）。"
+)
+
+if __name__ == "__main__":
+    iface.launch()