add: image generation and refactor prompt

app.py

@@ -1,10 +1,10 @@
 import os
 import tempfile
-
+import requests
 import streamlit as st
-
 from src.pipeline.main import LearnableAI
 from src.services.use_case_three.story_generator import StoryGenerator
+from src.services.image_generation.image_gen import image_generation
 
 
 def main():
@@ -47,7 +47,23 @@ def main():
 
                     st.success("Content Generated Successfully!")
                     st.write("### Results")
-                    st.write(result)
+                    st.write(result.split("</think>")[1].split("Image Generation Prompt")[0])
+                    image_prompt = result.split("Image Generation Prompt")[1]
+                    input = {
+                        "prompt": image_prompt,
+                        "go_fast": True,
+                        "megapixels": "1",
+                        "num_outputs": 1,
+                        "aspect_ratio": "1:1",
+                        "output_format": "webp",
+                        "output_quality": 80,
+                        "num_inference_steps": 4
+                    }
+                    output_image_url = image_generation(input)
+                    image_bytes = requests.get(output_image_url).content
+                    st.write("### Generated Image")
+                    st.image(image_bytes, caption="Generated Image")
+
 
                 except Exception as e:
                     st.error(f"An error occurred: {str(e)}")

src/services/image_generation/__init__.py

@@ -0,0 +1,5 @@
+"""
+project @ LearnableAI
+created @ 2025-01-17
+author  @ github.com/ishworrsubedii
+"""

src/services/image_generation/image_gen.py

@@ -0,0 +1,10 @@
+import replicate
+
+
+def image_generation(prompt: dict):
+    output = replicate.run(
+        "black-forest-labs/flux-schnell",
+        input=prompt,
+    )
+
+    return output[0]

src/services/use_case_one/word_to_sentence.py

@@ -89,13 +89,13 @@ class UseCaseOne:
             completion = self.client.chat.completions.create(
                 messages=[{
                     "role": "system",
-                    "content": "You are an expert in creating educational content for children with ADHD."
+                    "content": "You are an expert in creating educational content for children with ADHD. and you need to generate some educational content for children to understand the concept of the given words."
                 },
                     {
                         "role": "user",
                         "content": prompt
                     }],
-                model="mixtral-8x7b-32768",
+                model="deepseek-r1-distill-llama-70b",
                 temperature=1.0,  # if increase temperature, it will generate more creative content
                 max_tokens=2000,
                 top_p=1,

src/services/use_case_three/story_generator.py

@@ -9,7 +9,7 @@ class StoryGenerator:
     def process_image(self, image_path):
         try:
             img = Image.open(image_path)
-            prompt = self._create_story_prompt()
+            prompt = self.createstory_prompt()
 
             response = self.model.generate_content([prompt, img])
             return response.text
@@ -17,17 +17,17 @@ class StoryGenerator:
         except Exception as e:
             return f"Error processing image: {str(e)}"
 
-    def _create_story_prompt(self):
-        return """
-        Create a delightful and imaginative children's story inspired by the following captions from up to 4 images. Combine all captions into a single storyline that flows naturally. Use simple, child-friendly language and a playful tone suitable for children aged 4–8.
-        
-        Ensure the story has:
-        
-        A central character (e.g., a curious child, an animal, or a magical creature) that ties the events together.
-        A clear beginning, middle, and end with an engaging plot.
-        Elements of wonder and adventure to capture a child's imagination.
-        A gentle life lesson or moral, encouraging themes like kindness, courage, or curiosity.
-        Include vivid descriptions of settings and actions to match the imagery, and ensure the captions inspire key moments in the story.
-        
-        Please incorporate all these captions seamlessly into the storyline
-        """
+    def createstory_prompt(self):
+        return """You are a skilled children's story writer. Your task is to create a simple, engaging children's story that naturally incorporates the provided image captions into a cohesive narrative.
+                Context:
+                - Target audience: Young children
+                - Output format: Markdown file
+                - Story length: 4-5 sentences in line by line separated by \n
+                Image Captions to incorporate:
+                {captions}
+                Writing Guidelines:
+                - Use age-appropriate vocabulary
+                - Avoid complex sentence structures
+                - Create clear connections between events
+                - Keep the narrative positive and engaging
+                Your response should be formatted as a proper Markdown file"""