Upload 4 files

img_gen_v2.py

@@ -0,0 +1,72 @@
+import torch
+
+from diffusers import StableDiffusionImg2ImgPipeline, \
+    StableDiffusionPipeline
+
+
+def check_cuda_device():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    return device
+
+
+def get_the_model(device=None):
+    model_id = "stabilityai/stable-diffusion-2"
+    # if path:
+    #     pipe = StableDiffusionPipeline.from_pretrained(path, torch_dtype=torch.float16)
+    # else:
+    pipe = StableDiffusionPipeline.from_pretrained(model_id,
+                                                   torch_dtype=torch.float16)
+    if device:
+        pipe.to(device)
+    else:
+        device = check_cuda_device()
+        pipe.to(device)
+
+    return pipe
+
+
+def get_image_to_image_model(path=None, device=None):
+    model_id = "stabilityai/stable-diffusion-2"
+    if path:
+        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+            path,
+            torch_dtype=torch.float16)
+    else:
+        pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+            model_id,
+            torch_dtype=torch.float16)
+    if device:
+        if device == "cuda" or device == "cpu":
+            pipe.to(device)
+    else:
+        device = check_cuda_device()
+        pipe.to(device)
+
+    return pipe
+
+
+def gen_initial_img(int_prompt):
+    # image = get_the_model(num_inference_steps=100).images[0]
+    model = get_the_model(None)
+    image = model(int_prompt, num_inference_steps=100).images[0]
+
+    return image
+
+
+def generate_story(int_prompt, steps, iterations=100):
+    image_dic = {}
+    init_img = gen_initial_img(int_prompt)
+    img2img_model = get_image_to_image_model()
+
+    img = init_img
+
+    for idx, step in enumerate(steps):
+        image = img2img_model(prompt=step, image=img, strength=0.75, guidance_scale=7.5,
+                              num_inference_steps=iterations).images[0]
+        image_dic[idx] = {
+            "image": image,
+            "prompt": step
+        }
+        img = image
+
+    return init_img, image_dic

requirements.txt

@@ -3,3 +3,9 @@ langchain==0.0.153
 openai==0.27.5
 anthropic==0.2.7
 python-dotenv==1.0.0
+gTTS==2.3.2
+torch==2.0.0
+diffusers==0.16.1
+transformers
+ftfy
+accelerate

streamlit_app.py

@@ -0,0 +1,95 @@
+import streamlit as st
+import requests
+from PIL import Image
+from io import BytesIO
+from gtts import gTTS
+
+from img_gen_v2 import generate_story
+from prompt_generation import pipeline
+
+
+# Function to create the page navigation
+def page_navigation(current_page):
+    col1, col2, col3 = st.columns(3)
+
+    if current_page > 0:
+        with col1:
+            if st.button('<< Previous'):
+                current_page -= 1
+
+    with col2:
+        st.write(f'Step {current_page} of 10')
+
+    if current_page < 10:
+        with col3:
+            if st.button('Next >>'):
+                if current_page == 0:
+                    user_input = st.session_state.user_input
+                    prompt_response = pipeline(user_input, 10)
+                    steps = prompt_response.get("steps")
+                    init_prompt = prompt_response.get("story")
+
+                    init_img, img_dict = generate_story(init_prompt, steps)
+
+                    st.session_state.pipeline_response = prompt_response
+                    st.session_state.init_img = init_img
+                    st.session_state.img_dict = img_dict
+
+                current_page += 1
+
+    return current_page
+
+
+# Main function to display the pages
+def get_pipeline_data(page_number):
+    pipeline_response = st.session_state.pipeline_response
+    text_output = pipeline_response.get("steps")[page_number - 1]
+
+    # random_img = f"https://picsum.photos/800/600?random={page_number}"
+    # response = requests.get(random_img)
+    # image = Image.open(BytesIO(response.content))
+    img_dict = st.session_state.img_dict
+    img = img_dict[page_number-1]
+
+    return {"text_output": text_output, "image_obj": img}
+
+
+def main():
+    st.set_page_config(page_title="Narrative chat", layout="wide")
+    st.title("DreamBot")
+
+    # Initialize the current page
+    current_page = st.session_state.get('current_page', 0)
+
+    # Display content for each page
+    if current_page == 0:
+        st.write("Tell me what story you would like me to tell:")
+        user_input = st.text_area("")
+        st.session_state.user_input = user_input
+
+    else:
+        # Retrieve data from random generators
+        data = get_pipeline_data(current_page)
+        text_output = data.get('text_output', '')
+        image = data.get('image_obj', '')
+
+        # Display text output
+        st.write(text_output)
+
+        tts = gTTS(text_output)
+        tts.save('audio.mp3')
+        st.audio('audio.mp3')
+
+        # Display image output
+        if image:
+            st.image(image, use_column_width=False, width=400)
+
+    # Display page navigation
+    current_page = page_navigation(current_page)
+
+    st.write('current_page:', current_page)
+    st.session_state.current_page = current_page
+
+
+if __name__ == "__main__":
+    main()