Upload 2 files

app.py

@@ -1,12 +1,17 @@
 import gradio as gr
 import numpy as np
 from bark import SAMPLE_RATE, generate_audio
-from transformers import pipeline, set_seed
+from transformers import pipeline
+from diffusers import DiffusionPipeline
 import time
 import torch
+import random
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 generate_story_pipe = pipeline("text-generation", model="openai-community/gpt2")
-set_seed(42)
+image_pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype= torch.float16 if torch.cuda.is_available() else torch.float32)
+image_pipe = image_pipe.to(device)
 
 # Para medir el rendimiento de los métodos, voy a crear este decorador, que simplemente imprime en nuestra terminal el tiempo de ejecucion de los metodos que tengan los modelos y los usen, de esta manera podremos estudiar el tiempo que este cada modelo activo
 def measure_performance(func):
@@ -20,13 +25,6 @@ def measure_performance(func):
         return result
     return wrapper
 
-# Ejemplos que pondre debajo de la ventana
-examples = [
-    ["A love story"],
-    ["An humor story"],
-    ["A history about a vampire who dance Reggaeton"],
-]
-
 @measure_performance
 def generate_story(theme):
     prompt = f"Tell me a storie about {theme}"
@@ -40,6 +38,21 @@ def gen_tts(text):
     audio_arr = (audio_arr * 32767).astype(np.int16)
     return (SAMPLE_RATE, audio_arr)
 
+@measure_performance
+def generate_image(prompt):
+    seed = random.randint(0, 1000)
+    generator = torch.Generator().manual_seed(seed)
+
+    image = image_pipe(
+        prompt=prompt,
+        num_inference_steps=2,
+        width=256,
+        height=256,
+        generator=generator,
+    ).images[0]
+
+    return image
+
 def generate_story_and_convert_to_audio(theme):
     story_text = generate_story(theme)
     return gen_tts(story_text)
@@ -50,6 +63,13 @@ def subtitle():
                 <p>Cuando se necesite dormir a los mas pequeños de la casa, podrá usar este espacio, en el cual se integrarán varias inteligencias artificiales diseñadas para contar historias de todo tipo, al estilo del libro Las mil y una noches.</p>
     """)
 
+# Ejemplos que pondre debajo de la ventana
+examples = [
+    ["A love story"],
+    ["An humor story"],
+    ["A history about a vampire who dance Reggaeton"],
+]
+
 # Documentacion
 def end_text():
     gr.Markdown("<h2>Documentación</h2>")
@@ -123,10 +143,12 @@ with gr.Blocks(css=css) as block:
             story_theme = gr.Textbox(label="Escribe de que quieres que vaya la historai", placeholder="Introduce la descripción de la historia aquí...")
         with gr.Column():
             audio_out = gr.Audio(label="Generated Audio", type="numpy", elem_id="audio_out")
+            image_out = gr.Image(label="Generated image", show_label=False)
     run_button = gr.Button("Cuéntame la historia")
     run_button.click(fn=lambda: gr.update(visible=False), inputs=None, outputs=None, queue=False).then(
         fn=generate_story_and_convert_to_audio, inputs=[story_theme], outputs=[audio_out], queue=True).then(
-            fn=lambda: gr.update(visible=True), inputs=None, outputs=None, queue=False)
+        fn=lambda theme: generate_image(theme), inputs=[story_theme], outputs=[image_out], queue=True).then(
+        fn=lambda: gr.update(visible=True), inputs=None, outputs=None, queue=False)
 
 
     gr.Examples(examples=examples, fn=generate_story_and_convert_to_audio, inputs=[story_theme],

requirements.txt

@@ -2,4 +2,6 @@ torch
 transformers
 sentencepiece
 sacremoses
-bark
+bark
+diffusers
+random