add upload to r2

app.py

@@ -14,6 +14,8 @@ from huggingface_hub import hf_hub_download, HfFileSystem, ModelCard, snapshot_d
 import copy
 import random
 import time
+import boto3
+from io import BytesIO
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
@@ -49,6 +51,28 @@ class calculateDuration:
             print(f"Elapsed time: {self.elapsed_time:.6f} seconds")
 
 
+def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
+    print("upload_image_to_r2", 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
+
+
 @spaces.GPU(duration=70)
 def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, progress):
     pipe.to("cuda")
@@ -69,7 +93,7 @@ def generate_image(prompt, steps, seed, cfg_scale, width, height, lora_scale, pr
             yield img
 
 
-def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
     
     with calculateDuration("Unloading LoRA"):
         pipe.unload_lora_weights()
@@ -88,13 +112,20 @@ def run_lora(prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, see
     # Consume the generator to get the final image
     final_image = None
     step_counter = 0
+    final_image = None
     for image in image_generator:
         step_counter+=1
         final_image = image
         progress_bar = f'<div class="progress-container"><div class="progress-bar" style="--current: {step_counter}; --total: {steps};"></div></div>'
-        yield image, seed, gr.update(value=progress_bar, visible=True)
-        
-    yield final_image, seed, gr.update(value=progress_bar, visible=False)
+        yield image, seed, gr.update(value=progress_bar, visible=True),  json.dumps({"status": "processing"})
+    
+    if upload_to_r2:
+        url = upload_image_to_r2(final_image, account_id, access_key, secret_key, bucket)
+        result = {"status": "success", "url": url}
+    else:
+        result = {"status": "success", "message": "Image generated but not uploaded"}
+    
+    yield final_image, seed, gr.update(value=progress_bar, visible=False),  json.dumps(result)
 
 
 css="""
@@ -129,15 +160,23 @@ with gr.Blocks(css=css) as demo:
                     cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5)
                     steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28) 
 
+                upload_to_r2 = gr.Checkbox(label="Upload 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")
+        
+
         with gr.Column():
             progress_bar = gr.Markdown(elem_id="progress",visible=False)
             result = gr.Image(label="Result", show_label=False)
+            json_text = gr.Text()
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = run_lora,
-        inputs = [prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale],
-        outputs=[result, seed, progress_bar]
+        inputs = [prompt, cfg_scale, steps, lora_repo, lora_name, randomize_seed, seed, width, height, lora_scale, upload_to_r2, account_id, access_key, secret_key, bucket],
+        outputs=[result, seed, progress_bar, json_text]
     )
 
 demo.queue().launch()

requirements.txt

@@ -3,4 +3,5 @@ git+https://github.com/huggingface/diffusers@3b604e8c384631e1f66a4fd9076ed5e7e2b
 spaces
 transformers
 peft
-sentencepiece
+sentencepiece
+boto3