Added S3 integration

app.py

@@ -4,9 +4,21 @@ from io import BytesIO
 from PIL import Image
 import gradio as gr
 import os
+import boto3
 
 model_key = os.environ.get("model_key")
 api_key = os.environ.get("api_key")
+aws_access_key_id = os.environ.get("aws_access_key_id")
+aws_secret_access_key = os.environ.get("aws_secret_access_key")
+
+# Create a session using AWS credentials
+session = boto3.Session(aws_access_key_id, aws_secret_access_key)
+
+# Create an S3 resource object using the session
+s3 = session.resource('s3')
+
+# Select your bucket
+bucket = s3.Bucket('bwlmonet')
 
 model_inputs = {
   "endpoint": "txt2img",
@@ -25,6 +37,9 @@ model_inputs = {
   }
 }
 
+for obj in bucket.objects.all():
+  print(obj.key)
+
 def stable_diffusion_txt2img(prompt, api_key, model_key, model_inputs):
   # Update the model_inputs with the provided prompt
   model_inputs["params"]["prompt"] = prompt
@@ -37,6 +52,16 @@ def stable_diffusion_txt2img(prompt, api_key, model_key, model_inputs):
   image_encoded = image_byte_string[0].encode("utf-8")
   image_bytes = BytesIO(base64.b64decode(image_encoded))
   image = Image.open(image_bytes)
+
+  # Save image to S3
+  key = f"{prompt}.png"
+  image.save(key)
+  with open(key, "rb") as data:
+    bucket.put_object(Key=key, Body=data)
+  
+  for obj in bucket.objects.all():
+    print(obj.key)
+
   return image
 
 # Gradio Interface