jiuface commited on
Commit
0a8ba8d
·
1 Parent(s): 91d2029
Files changed (1) hide show
  1. app.py +19 -14
app.py CHANGED
@@ -84,25 +84,30 @@ def get_depth_map(image):
84
  return image
85
 
86
 
87
- def upload_to_s3(image, region, access_key, secret_key, bucket_name):
 
 
 
88
  s3 = boto3.client(
89
  's3',
90
- region_name=region,
 
91
  aws_access_key_id=access_key,
92
  aws_secret_access_key=secret_key
93
  )
94
- image_key = f"generated_images/{random.randint(0, MAX_SEED)}.png"
95
  buffer = BytesIO()
96
  image.save(buffer, "PNG")
97
  buffer.seek(0)
98
-
99
- s3.upload_fileobj(buffer, bucket_name, image_key)
100
- return image_key
 
101
 
102
 
103
 
104
  @spaces.GPU(enable_queue=True)
105
- 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()):
106
 
107
  if image_url:
108
  orginal_image = load_image(image_url)
@@ -126,7 +131,7 @@ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, contr
126
  ).images[0]
127
 
128
  if upload_to_s3:
129
- url = upload_to_s3(generated_image, region, access_key, secret_key, bucket)
130
  result = {"status": "success", "url": url}
131
  else:
132
  result = {"status": "success", "message": "Image generated but not uploaded"}
@@ -154,11 +159,11 @@ with gr.Blocks() as demo:
154
  value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
155
  )
156
 
157
- upload_to_s3 = gr.Checkbox(label="Upload to S3", value=False)
158
- region = gr.Textbox(label="S3 Region", placeholder="Enter S3 region here")
159
- access_key = gr.Textbox(label="Access Key", placeholder="Enter S3 access key here")
160
- secret_key = gr.Textbox(label="Secret Key", placeholder="Enter S3 secret key here")
161
- bucket = gr.Textbox(label="Bucket Name", placeholder="Enter S3 bucket name here")
162
 
163
 
164
  with gr.Column():
@@ -175,7 +180,7 @@ with gr.Blocks() as demo:
175
  control_strength,
176
  seed,
177
  upload_to_s3,
178
- region,
179
  access_key,
180
  secret_key,
181
  bucket
 
84
  return image
85
 
86
 
87
+ def upload_to_s3(image, account_id, access_key, secret_key, bucket_name):
88
+
89
+ connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
90
+
91
  s3 = boto3.client(
92
  's3',
93
+ endpoint_url=connectionUrl,
94
+ region_name='auto',
95
  aws_access_key_id=access_key,
96
  aws_secret_access_key=secret_key
97
  )
98
+ image_file = f"generated_images/{random.randint(0, MAX_SEED)}.png"
99
  buffer = BytesIO()
100
  image.save(buffer, "PNG")
101
  buffer.seek(0)
102
+ print("start to upload")
103
+ s3.upload_fileobj(buffer, bucket_name, image_file)
104
+ print("upload finish", image_file)
105
+ return image_file
106
 
107
 
108
 
109
  @spaces.GPU(enable_queue=True)
110
+ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, control_strength, seed, upload_to_s3, account_id, access_key, secret_key, progress=gr.Progress()):
111
 
112
  if image_url:
113
  orginal_image = load_image(image_url)
 
131
  ).images[0]
132
 
133
  if upload_to_s3:
134
+ url = upload_to_s3(generated_image, account_id, access_key, secret_key, bucket)
135
  result = {"status": "success", "url": url}
136
  else:
137
  result = {"status": "success", "message": "Image generated but not uploaded"}
 
159
  value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
160
  )
161
 
162
+ upload_to_s3 = gr.Checkbox(label="Upload to R2", value=False)
163
+ account_id = gr.Textbox(label="Account Id", placeholder="Enter R2 account id")
164
+ access_key = gr.Textbox(label="Access Key", placeholder="Enter R2 access key here")
165
+ secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
166
+ bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here")
167
 
168
 
169
  with gr.Column():
 
180
  control_strength,
181
  seed,
182
  upload_to_s3,
183
+ account_id,
184
  access_key,
185
  secret_key,
186
  bucket