Update app.py

app.py

@@ -8,6 +8,33 @@ import gradio as gr
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
 from peft import PeftModel
+from openai import OpenAI
+from dotenv import load_dotenv
+
+load_dotenv()
+
+openai_key = os.getenv("OPENAI_KEY")
+
+if openai_key == "":
+    sys.exit("Please Provide Your OpenAI API Key")
+
+
+def generate_image(text):
+    try:
+        client = OpenAI(api_key=openai_key)
+
+        response = client.images.generate(
+            model="dall-e-3",
+            prompt="Create an illustration that accurately depicts the character and the setting of this story:"+text,
+            n=1,
+            size="1024x1024"
+        )
+    except Exception as error:
+        print(str(error))
+        raise gr.Error("An error occurred while generating speech. Please check your API key and come back try again.")
+
+    return response.data[0].url
+    
 
 # Constants
 MAX_MAX_NEW_TOKENS = 2048
@@ -124,16 +151,19 @@ def generate(
         output = "".join(outputs)
         yield output
 
-    # final_story = "".join(outputs)
+    final_story = "".join(outputs)
+    generate_image(final_story)
     # try:
     #     saved_story = Story(message=message, content=final_story).save()
     #     yield f"{final_story}\n\n Story saved with ID: {saved_story.story_id}"
     # except Exception as e:
     #     yield f"Failed to save story: {str(e)}"
+    
 
 # Gradio Interface Setup
 chat_interface = gr.ChatInterface(
     fn=generate,
+    label=" ",
     fill_height=True,
     stop_btn=None,
     examples=[
@@ -146,7 +176,7 @@ chat_interface = gr.ChatInterface(
 
 # Gradio Web Interface
 with gr.Blocks(css=custom_css,theme='shivi/calm_seafoam',fill_height=True) as demo:
-    out = gr.Textbox(placeholder="What is your name?")
+    output_image = gr.Image(label="Image Output")
     chat_interface.render()
     # gr.Markdown(LICENSE)