Update src/app.py

src/app.py

@@ -15,6 +15,9 @@ from refiners.foundationals.latent_diffusion import Solver, solvers
 import requests
 from enhancer import ESRGANUpscaler, ESRGANUpscalerCheckpoints
 import time
+import boto3
+from datetime import datetime
+import json
 
 pillow_heif.register_heif_opener()
 pillow_heif.register_avif_opener()
@@ -101,6 +104,28 @@ enhancer.to(device=DEVICE, dtype=DTYPE)
 
 
 
+def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
+    print("upload_image_to_s3", account_id, access_key, secret_key, bucket_name)
+    connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
+
+    s3 = boto3.client(
+        's3',
+        endpoint_url=connectionUrl,
+        region_name='auto',
+        aws_access_key_id=access_key,
+        aws_secret_access_key=secret_key
+    )
+
+    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+    image_file = f"generated_images/{current_time}_{random.randint(0, MAX_SEED)}.png"
+    buffer = BytesIO()
+    image.save(buffer, "PNG")
+    buffer.seek(0)
+    s3.upload_fileobj(buffer, bucket_name, image_file)
+    print("upload finish", image_file)
+    return image_file
+    
+
 class calculateDuration:
     def __init__(self, activity_name=""):
         self.activity_name = activity_name
@@ -134,7 +159,12 @@ def process(
     denoise_strength: float = 0.35,
     num_inference_steps: int = 18,
     solver: str = "DDIM",
-) -> tuple[Image.Image, Image.Image]:
+    upload_to_r2: bool = True, 
+    account_id = "", 
+    access_key = "", 
+    secret_key = "", 
+    bucket = ""
+) -> Tuple[Tuple[Image.Image, Image.Image], str]:
     manual_seed(seed)
     
     if image_url:
@@ -164,7 +194,14 @@ def process(
             solver_type=solver_type,
         )
         print("enhancer finish")
-    return [input_image, enhanced_image]
+    
+    if upload_to_r2:
+        url = upload_image_to_r2(enhanced_image, account_id, access_key, secret_key, bucket_name)
+        result = {"status": "success", "url": url}
+    else:
+        result = {"status": "success", "message": "Image generated but not uploaded"}
+    return [input_image, enhanced_image], json.dumps(result)
+    
 
 with gr.Blocks() as demo:
     gr.HTML(TITLE)
@@ -176,7 +213,10 @@ with gr.Blocks() as demo:
             run_button = gr.ClearButton(components=None, value="Enhance Image")
         with gr.Column():
             output_slider = ImageSlider(label="Generate image", type="pil", slider_color="pink")
+            logs = gr.Textbox(label="logs")
+            
             run_button.add(output_slider)
+            
 
     with gr.Accordion("Advanced Options", open=False):
         prompt = gr.Textbox(
@@ -256,6 +296,13 @@ with gr.Blocks() as demo:
             label="Solver",
         )
 
+        upload_to_r2 = gr.Checkbox(label="Upload generated image to R2", value=False)
+        account_id = gr.Textbox(label="Account Id", placeholder="Enter R2 account id")
+        access_key = gr.Textbox(label="Access Key", placeholder="Enter R2 access key here")
+        secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
+        bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here")
+        
+
     run_button.click(
         fn=process,
         inputs=[
@@ -273,8 +320,13 @@ with gr.Blocks() as demo:
             denoise_strength,
             num_inference_steps,
             solver,
+            upload_to_r2,
+            account_id,
+            access_key,
+            secret_key,
+            bucket
         ],
-        outputs=output_slider,
+        outputs=[output_slider, logs]
     )
 
 demo.launch(share=False)