Upload 3 files

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 from gtts import gTTS
 
-from img_gen_v2 import generate_story
+from img_gen import generate_story
 from prompt_generation import pipeline
 
 
@@ -15,7 +15,7 @@ def page_navigation(current_page):
                 current_page -= 1
 
     with col2:
-        st.write(f'Step {current_page} of 10')
+        print(f'Step {current_page} of 10')
 
     if current_page < 10:
         with col3:
@@ -23,10 +23,11 @@ def page_navigation(current_page):
                 if current_page == 0:
                     user_input = st.session_state.user_input
                     prompt_response = pipeline(user_input, 10)
-                    steps = prompt_response.get("steps")
+                    image_prompts_steps = prompt_response.get("image_prompts_steps")
                     init_prompt = prompt_response.get("story")
 
-                    init_img, img_dict = generate_story(init_prompt, steps)
+                    init_img, img_dict = generate_story(init_prompt,
+                                                        image_prompts_steps)
 
                     st.session_state.pipeline_response = prompt_response
                     st.session_state.init_img = init_img
@@ -42,7 +43,7 @@ def get_pipeline_data(page_number):
     pipeline_response = st.session_state.pipeline_response
     text_output = pipeline_response.get("steps")[page_number - 1]
     img_dict = st.session_state.img_dict
-    img = img_dict[page_number-1].get("image")
+    img = img_dict[page_number - 1].get("image")
 
     return {"text_output": text_output, "image_obj": img}
 
@@ -56,7 +57,7 @@ def main():
 
     # Display content for each page
     if current_page == 0:
-        st.write("Tell me what story you would like me to tell:")
+        st.write("Describe a story you would like me to tell:")
         user_input = st.text_area("")
         st.session_state.user_input = user_input
 
@@ -69,7 +70,7 @@ def main():
         # Display text output
         st.write(text_output)
 
-        tts = gTTS(text_output)
+        tts = gTTS(text_output.split(".", 1)[1])
         tts.save('audio.mp3')
         st.audio('audio.mp3')
 

img_gen.py

@@ -0,0 +1,69 @@
+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"
+    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):
+    model = get_the_model(None)
+    image = model(int_prompt, num_inference_steps=100).images[0]
+
+    return image
+
+
+def generate_story(int_prompt, steps, iterations=133):
+    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):
+        print(f"step: {idx}")
+        print(step)
+        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

prompt_generation.py

@@ -97,6 +97,7 @@ def pipeline(user_description: str, n_steps: int = 10) -> dict:
 
     image_prompts = [fut.result() for fut in image_prompts_futures]
 
+    print(story)
     return {"story": story, "steps": steps, "image_prompts": image_prompts}