upload img to s3

app.py

@@ -17,6 +17,8 @@ from transformers import DPTFeatureExtractor, DPTForDepthEstimation, DPTImagePro
 from transformers import CLIPImageProcessor
 from diffusers.utils import load_image
 from gradio_imageslider import ImageSlider
+import boto3
+from io import BytesIO
 
 device = "cuda"
 base_model_id = "SG161222/RealVisXL_V4.0"
@@ -82,9 +84,25 @@ def get_depth_map(image):
     return image
 
 
+def upload_to_s3(image, region, access_key, secret_key, bucket_name):
+    s3 = boto3.client(
+        's3',
+        region_name=region,
+        aws_access_key_id=access_key,
+        aws_secret_access_key=secret_key
+    )
+    image_key = f"generated_images/{random.randint(0, MAX_SEED)}.png"
+    buffer = BytesIO()
+    image.save(buffer, "PNG")
+    buffer.seek(0)
+    
+    s3.upload_fileobj(buffer, bucket_name, image_key)
+    return image_key
+
+
 
 @spaces.GPU(enable_queue=True)
-def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, control_strength, seed, progress=gr.Progress()):
+def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, control_strength, seed, upload_to_s3, region, access_key, secret_key, progress=gr.Progress()):
 
     if image_url:
         orginal_image = load_image(image_url)
@@ -106,7 +124,14 @@ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, contr
         generator=generator,
         image=depth_image
     ).images[0]
-    return [[depth_image, generated_image], "ok"]
+
+    if upload_to_s3:
+        url = upload_to_s3(generated_image, region, access_key, secret_key, bucket)
+        result = {"status": "success", "url": url}
+    else:
+        result = {"status": "success", "message": "Image generated but not uploaded"}
+    
+    return [[depth_image, generated_image], json.dumps(result)]
 
 with gr.Blocks() as demo:
     
@@ -128,6 +153,14 @@ with gr.Blocks() as demo:
                     label="Negative prompt",
                     value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
                 )
+
+                upload_to_s3 = gr.Checkbox(label="Upload to S3", value=False)
+                region = gr.Textbox(label="S3 Region", placeholder="Enter S3 region here")
+                access_key = gr.Textbox(label="Access Key", placeholder="Enter S3 access key here")
+                secret_key = gr.Textbox(label="Secret Key", placeholder="Enter S3 secret key here")
+                bucket = gr.Textbox(label="Bucket Name", placeholder="Enter S3 bucket name here")
+        
+
         with gr.Column():
             result = ImageSlider(label="Generate image", type="pil", slider_color="pink")
             logs = gr.Textbox(label="logs")
@@ -140,7 +173,12 @@ with gr.Blocks() as demo:
         num_steps,
         guidance_scale,
         control_strength,
-        seed
+        seed,
+        upload_to_s3,
+        region,
+        access_key,
+        secret_key,
+        bucket
     ]
     run_button.click(
             fn=randomize_seed_fn,

requirements.txt

@@ -9,4 +9,5 @@ requests
 spaces
 huggingface_hub
 controlnet-aux
-safetensors
+safetensors
+boto3