Implememted Image Generation api

main.py

@@ -5,11 +5,11 @@ from my_agent.agent import build_graph
 import pandas as pd
 from typing import Optional , List
 from my_agent.utils.initial_interaction import BusinessInteractionChatbot
-from my_agent.utils.utils import encode_image_to_base64 , generate_final_story 
+from my_agent.utils.utils import encode_image_to_base64 , generate_final_story, generate_image
+
 
 import json
 from my_agent.utils.nodes import generate_final_story
-from collections import defaultdict
 
 # Store brainstorming results per thread_id
 
@@ -33,9 +33,6 @@ def business_chat(msg: UserMessage):
 
 
 
-
-
-
 @app.post("/brainstrom")
 def run_graph(
     query: List[str],  # sent as JSON body
@@ -51,7 +48,7 @@ def run_graph(
         'topic': query,
         'images': image_base64_list,
         'latest_preferred_topics':preferred_topics,
-        'business_details': stored_data['business_details']
+        # 'business_details': stored_data['business_details']
     },
     config={"configurable": {"thread_id": thread_id}})
     stored_data['brainstroming_response']=result
@@ -62,16 +59,24 @@ def run_graph(
     }
 
 
-
-
 @app.post("/generate-final-story")
-def generate_final_story_endpoint(
-):
+def generate_final_story_endpoint():
     final_story =   generate_final_story(stored_data["brainstroming_response"])
+    stored_data['final_story']=final_story
     return {
         'response': final_story
     }
-    
+
+@app.post("/generate-image")  
+def generate_image_endpoint():
+    image = generate_image(str(stored_data['final_story']))
+    stored_data['generated_image']=image
+    return {
+        'image':image
+    }
+
+
+
 
 
 

my_agent/utils/prompts.py

@@ -0,0 +1,11 @@
+story_to_prompt = '''
+You are an expert in crafting vivid, cinematic prompts for image generation models like Midjourney, DALL·E, or Stable Diffusion.
+
+Given a paragraph describing a visual scene, rewrite it as a single, detailed sentence that captures the atmosphere, key visual elements, emotions, and composition — as if directing a high-quality cinematic image.
+
+Example input:
+"A bustling street scene sets the tone for a culinary journey, as the camera pans across the exterior of a cozy vegetarian restaurant, with the sign 'Vegetarian Delights' prominently displayed. The camera takes us inside, where the warm and inviting atmosphere is filled with the aroma of spices and the soft glow of lighting, highlighting the menu and the chefs expertly preparing a variety of vegetarian dishes. A satisfied customer takes a bite of their meal, with a look of delight on their face, as the camera zooms in on the flavors and textures of the dish, evoking a sense of joy and contentment."
+
+Desired output:
+"A cinematic street scene reveals the charming exterior of 'Vegetarian Delights,' then glides into a warmly lit, spice-scented interior where skilled chefs craft vibrant vegetarian dishes and a delighted diner savors each flavorful bite in glowing ambiance."
+'''

my_agent/utils/utils.py

@@ -6,6 +6,9 @@ import base64
 from PIL import Image
 from io import BytesIO
 from fastapi import UploadFile
+from huggingface_hub import InferenceClient
+from .prompts import story_to_prompt
+
 def generate_final_story(query):
     if len(query['preferred_topics'])>0:
         template = f'''I want to create a detailed storyline for a video in the given topic. You have to provide me that storyline what to include in the video.
@@ -39,4 +42,47 @@ def encode_image_to_base64(uploaded_file: UploadFile) -> str:
 # Convert base64 string to PIL image (optional for LangGraph processing)
 def process_image(base64_str: str) -> Image.Image:
     image_data = base64.b64decode(base64_str)
-    return Image.open(BytesIO(image_data))
+    return Image.open(BytesIO(image_data))
+
+
+def generate_prompt(final_story):
+    print('************Entering prompt generator****************')
+    messages = [
+    (
+        "system",
+        story_to_prompt,
+    ),
+    ("human", final_story),
+]
+    prompt = llm.invoke(messages)
+    print('The prompt is:',prompt)
+    return prompt.content
+
+def generate_image(final_story):
+    prompt = generate_prompt(final_story)
+    print('************Finished prompt generator****************')
+
+    client = InferenceClient(
+        provider="hf-inference",
+        api_key="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",
+    )
+
+    print('************Finished calling generator****************')
+
+
+        # output is a PIL.Image object
+    image = client.text_to_image(
+        prompt,
+        model="black-forest-labs/FLUX.1-schnell",
+    )
+    print('*****************Image Created*******************')
+
+    image.save('image.png')
+    print('*****************Image Saved*******************')
+    return "Image Created"
+    # try:
+    #     return image
+    # except:
+    #     return 'Image created'
+
+