Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import gradio as gr
 import torch
 import nltk
-import random
 from openai import OpenAI
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from diffusers import StableDiffusionPipeline
@@ -10,6 +9,8 @@ from ultralytics import YOLO
 from gtts import gTTS
 from PIL import Image
 import numpy as np
+from nltk.tokenize import sent_tokenize
+from IPython.display import Audio
 
 # Βεβαιωθείτε ότι το API Key υπάρχει
 api_key = os.getenv("OPENAI_API_KEY")
@@ -20,57 +21,79 @@ if not api_key:
 client = OpenAI(api_key=api_key)
 
 # Φόρτωση μοντέλων
-print("Loading models...")
-yolo_model = YOLO("yolov8s.pt")  # Μοντέλο Object Detection
-text_generation = pipeline("text-generation", model="gpt2")
+yolo_model = YOLO("yolov8s.pt")
 stable_diffusion = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
 nltk.download("punkt")
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
-def generate_story(prompt):
-    response = client.completions.create(
-        model="gpt-4o",
-        prompt=prompt,
-        max_tokens=300
+def detect_objects(image_path):
+    results = yolo_model(image_path)
+    detected_objects = []
+    for r in results:
+        for box in r.boxes:
+            class_id = int(box.cls[0])
+            label = yolo_model.names[class_id]
+            detected_objects.append(label)
+    return detected_objects
+
+def generate_story(detected_objects):
+    story_prompt = f"Write a short story based on the following objects: {', '.join(detected_objects)}"
+    response = client.chat.completions.create(
+        model="gpt-4o-mini",
+        messages=[{"role": "user", "content": story_prompt}],
+        max_tokens=200
     )
-    return response.choices[0].text.strip()
+    return response.choices[0].message.content
 
-def detect_objects(image):
-    results = yolo_model(image)
-    return results[0].plot()
+def summarize_story(story):
+    summary = summarizer(story, max_length=100, do_sample=False)[0]['summary_text']
+    scenes = sent_tokenize(summary)
+    return scenes
 
-def generate_image(prompt):
-    image = stable_diffusion(prompt).images[0]
-    return image
+def generate_images(story):
+    scenes = summarize_story(story)
+    prompts = [f"Highly detailed, cinematic scene: {scene}, digital art, 4K, realistic lighting" for scene in scenes]
+    images = []
+    for prompt in prompts:
+        image = stable_diffusion(prompt).images[0]
+        images.append(image)
+    return images
 
-def text_to_speech(text):
-    tts = gTTS(text=text, lang="en")
-    tts.save("output.mp3")
-    return "output.mp3"
+def text_to_speech(story):
+    tts = gTTS(text=story, lang="en", slow=False)
+    audio_file_path = "story_audio.mp3"
+    tts.save(audio_file_path)
+    return audio_file_path
 
 demo = gr.Interface(
     fn={
-        "Generate Story": generate_story,
         "Detect Objects": detect_objects,
-        "Generate Image": generate_image,
+        "Generate Story": generate_story,
+        "Summarize Story": summarize_story,
+        "Generate Images": generate_images,
         "Text to Speech": text_to_speech,
     },
     inputs={
-        "Generate Story": gr.Textbox(placeholder="Write a story prompt..."),
         "Detect Objects": gr.Image(type="numpy"),
-        "Generate Image": gr.Textbox(placeholder="Describe an image..."),
+        "Generate Story": gr.Textbox(placeholder="Write a story prompt..."),
+        "Summarize Story": gr.Textbox(placeholder="Enter the generated story..."),
+        "Generate Images": gr.Textbox(placeholder="Enter the story for image generation..."),
         "Text to Speech": gr.Textbox(placeholder="Enter text to convert to speech...")
     },
     outputs={
+        "Detect Objects": "text",
         "Generate Story": "text",
-        "Detect Objects": "image",
-        "Generate Image": "image",
+        "Summarize Story": "text",
+        "Generate Images": "image",
         "Text to Speech": "audio"
     },
     title="AI-Powered Storytelling Assistant",
-    description="An AI assistant combining NLP, Object Detection, Image Generation, and TTS!"
+    description="An AI assistant combining Object Detection, NLP, Image Generation, and TTS!"
 )
 
 if __name__ == "__main__":
     demo.launch()
 
 
+
+