Add files via upload

app.py

@@ -0,0 +1,110 @@
+import streamlit as st
+import torch
+from PIL import Image
+import os
+from dotenv import load_dotenv
+import google.generativeai as genai
+from diffusers import StableDiffusionPipeline
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Configure Google Generative AI API
+api_key = os.getenv("GOOGLE_API_KEY")
+if api_key:
+    genai.configure(api_key=api_key)
+else:
+    st.error("GOOGLE_API_KEY is not set. Please set it in your environment.")
+    st.stop()
+
+@st.cache_resource
+def load_text_model():
+    model = genai.GenerativeModel("gemini-pro")
+    return model
+
+def generate_story_extension(user_input, model, temperature=0.85, max_output_tokens=1000):
+    generation_config = genai.types.GenerationConfig(
+        temperature=temperature,
+        max_output_tokens=max_output_tokens,
+    )
+    response = model.generate_content(user_input, generation_config=generation_config)
+    
+    if response.parts:
+        return ''.join(part.text for part in response.parts if hasattr(part, 'text'))
+    else:
+        return "No content generated."
+
+@st.cache_resource
+def load_image_model():
+    model_id = "stabilityai/stable-diffusion-2"
+    
+    try:
+        model = StableDiffusionPipeline.from_pretrained(
+            model_id, 
+            torch_dtype=torch.float32,
+            revision="fp16"
+        )
+        st.success(f"Successfully loaded the image generation model")
+        return model
+    except Exception as e:
+        st.error(f"An unexpected error occurred while loading the model: {str(e)}")
+        return None
+
+def main():
+    st.title("Interactive Storyteller")
+
+    text_model = load_text_model()
+    image_model = load_image_model()
+
+    if image_model is None:
+        st.error("Failed to load the image generation model. Please check the errors above and try again.")
+        return
+    
+    characters = st.text_area("Enter the name of characters along with their descriptions")
+    plot = st.text_area("Describe the plot in brief")
+    theme = st.text_area("Provide a theme of the story along with the setting")
+
+    if st.button("Generate Story"):
+        user_input = f"Theme and setting: {theme} plot: {plot} characters: {characters}"
+        prompt = f"Generate a complete story of at least 800 words based on input given by the user: \n {user_input}"
+        
+        with st.spinner("Generating initial story..."):
+            base_story = generate_story_extension(prompt, text_model)
+        
+        st.subheader("Generated Story")
+        st.write(base_story)
+
+        st.session_state.story = base_story
+        st.session_state.story_generated = True
+
+    if 'story_generated' in st.session_state and st.session_state.story_generated:
+        user_changed_input = st.text_input("If you want changes, type out the changes you want in the form of a simple prompt. Otherwise, leave blank to proceed.")
+        
+        if user_changed_input:
+            with st.spinner("Updating story..."):
+                prompt_change = f"Generate a new story from scratch of at least 800 words with reference to the previously generated story and based on changes instructed by the user: \n {user_changed_input}"
+                updated_story = generate_story_extension(prompt_change, text_model)
+            st.session_state.story = updated_story
+            st.subheader("Updated Story")
+            st.write(updated_story)
+
+        if st.button("Generate Images"):
+            with st.spinner("Generating story sequence and images..."):
+                prompt_for_img_gen = f"Based on the story: {st.session_state.story}, add \n delimiters to separate the story into at least 10 pivotal parts (each part represents a different pivotal chapter of the story) and at max 20 parts, where each pivotal part gives an illustrative description of 30 to 40 words about that part such that an image can be generated from the part and fed into a text-to-image model to show progressive story. Try involving new settings or new characters in each part."
+                
+                prompt_corpus = generate_story_extension(prompt_for_img_gen, text_model)
+                
+                story_sequence = [x for x in prompt_corpus.split("\n") if 'Part' in x]
+
+                st.subheader("Story Sequence")
+                for part in story_sequence:
+                    st.write(part)
+
+                st.subheader("Generated Images")
+                for i, part in enumerate(story_sequence):
+                    image = image_model(part).images[0]
+                    st.image(image, caption=f"Part {i+1}")
+                    st.write(part)
+
+if __name__ == "__main__":
+    main()