Update app.py

app.py

@@ -1,6 +1,11 @@
 import gradio as gr
 import torch
 from diffusers import StableDiffusionPipeline
+import os
+
+AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
+AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
+S3_BUCKET_NAME = os.getenv("BUCKET_NAME")
 
 model_id = "CompVis/stable-diffusion-v1-4"
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -10,15 +15,42 @@ pipe = StableDiffusionPipeline.from_pretrained(
 
 pipe = pipe.to(device)
 
-def generatorImage(name):
-    prompt = name
-    image = pipe(prompt).images[0]
-    image_name = '-'.join(prompt.split()) + ".png"
-    # image.save("./images/" + image_name)
-
-    return image_name
-
-
-
-iface = gr.Interface(fn=generatorImage, inputs="text", outputs="text")
+def text_to_image(text):
+    # Crea una instancia del cliente de S3
+    s3 = boto3.client('s3',
+                  aws_access_key_id=AWS_ACCESS_KEY_ID,
+                  aws_secret_access_key=AWS_SECRET_ACCESS_KEY)
+
+    def save_image_to_s3(image, image_name):
+        # Crea un objeto de BytesIO para almacenar la imagen
+        image_buffer = BytesIO()
+        image.save(image_buffer, format='PNG')
+        image_buffer.seek(0)
+    
+        # Define la ruta de destino en el bucket de S3
+        s3_key = f"public/{image_name}"
+    
+        # Sube la imagen al bucket de S3 en la ruta especificada
+        s3.upload_fileobj(image_buffer, bucket_name, s3_key)
+    
+    def generator_image(text):
+        prompt = text
+        image = pipe(prompt).images[0]
+        image_name = '-'.join(prompt.split()) + ".png"
+        
+        # Guarda la imagen en S3
+        save_image_to_s3(image, image_name)
+    
+        return image_name
+    
+    def generator_image_interface(text):
+        image_name = generator_image(text)
+        return f"Imagen generada: {image_name}"
+
+    # generate image
+    generator_image_interface(text);
+
+
+
+iface = gr.Interface(fn=text_to_image, inputs="text", outputs="text")
 iface.launch()