Update app.py

app.py

@@ -84,25 +84,41 @@ def generate_image(prompt, adapter_names, steps, seed, cfg_scale, width, height,
 
 
 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)
+    with calculateDuration("Upload images"):
+        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)
+        # start to generate thumbnail
+        thumbnail = image.copy()
+        thumbnail_width = 256
+        aspect_ratio = image.height / image.width
+        thumbnail_height = int(thumbnail_width * aspect_ratio)
+        thumbnail = thumbnail.resize((thumbnail_width, thumbnail_height), Image.LANCZOS)
+        
+        # Generate the thumbnail image filename
+        thumbnail_file = image_file.replace(".png", "_thumbnail.png")
+        
+        # Save thumbnail to buffer and upload
+        thumbnail_buffer = BytesIO()
+        thumbnail.save(thumbnail_buffer, "PNG")
+        thumbnail_buffer.seek(0)
+        s3.upload_fileobj(thumbnail_buffer, bucket_name, thumbnail_file)
+        print("upload thumbnail finish", thumbnail_file)
     return image_file
 
 def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed, width, height, upload_to_r2, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
@@ -164,9 +180,8 @@ def run_lora(prompt, lora_strings_json,  cfg_scale, steps, randomize_seed, seed,
         
     if final_image:
         if upload_to_r2:
-            with calculateDuration("Upload image"):
-                url = upload_image_to_r2(final_image, account_id, access_key, secret_key, bucket)
-                result = {"status": "success", "message": "upload image success", "url": url}
+            url = upload_image_to_r2(final_image, account_id, access_key, secret_key, bucket)
+            result = {"status": "success", "message": "upload image success", "url": url}    
         else:
             result = {"status": "success", "message": "Image generated but not uploaded"}
     else: