jiuface commited on
Commit
d792bfc
·
verified ·
1 Parent(s): 5f87b7d
Files changed (1) hide show
  1. app.py +7 -7
app.py CHANGED
@@ -66,7 +66,7 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
66
  return seed
67
 
68
 
69
- def get_depth_map(image):
70
  original_size = (image.size[1], image.size[0])
71
  image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
72
  with torch.no_grad(), torch.autocast("cuda"):
@@ -86,7 +86,7 @@ def get_depth_map(image):
86
  return image
87
 
88
 
89
- def upload_image_to_s3(image, account_id, access_key, secret_key, bucket_name):
90
  print("upload_image_to_s3", account_id, access_key, secret_key, bucket_name)
91
  connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
92
 
@@ -110,7 +110,7 @@ def upload_image_to_s3(image, account_id, access_key, secret_key, bucket_name):
110
 
111
 
112
  @spaces.GPU(enable_queue=True)
113
- def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, control_strength, seed, upload_to_s3, account_id, access_key, secret_key, bucket):
114
  print("process start")
115
  if image_url:
116
  print(image_url)
@@ -120,9 +120,9 @@ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, contr
120
 
121
  size = (orginal_image.size[0], orginal_image.size[1])
122
  print("image size", size)
123
- depth_image = get_depth_map(orginal_image)
124
  generator = torch.Generator().manual_seed(seed)
125
- print(prompt, n_prompt, guidance_scale, num_steps, control_strength, )
126
  generated_image = pipe(
127
  prompt=prompt,
128
  negative_prompt=n_prompt,
@@ -136,12 +136,12 @@ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, contr
136
  ).images[0]
137
 
138
  if upload_to_s3:
139
- url = upload_image_to_s3(generated_image, account_id, access_key, secret_key, bucket)
140
  result = {"status": "success", "url": url}
141
  else:
142
  result = {"status": "success", "message": "Image generated but not uploaded"}
143
 
144
- return [[orginal_image, generated_image], json.dumps(result)]
145
 
146
  with gr.Blocks() as demo:
147
 
 
66
  return seed
67
 
68
 
69
+ def get_depth_map(image, progress):
70
  original_size = (image.size[1], image.size[0])
71
  image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
72
  with torch.no_grad(), torch.autocast("cuda"):
 
86
  return image
87
 
88
 
89
+ def upload_image_to_s3(image, account_id, access_key, secret_key, bucket_name, progress):
90
  print("upload_image_to_s3", account_id, access_key, secret_key, bucket_name)
91
  connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
92
 
 
110
 
111
 
112
  @spaces.GPU(enable_queue=True)
113
+ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, control_strength, seed, upload_to_s3, account_id, access_key, secret_key, bucket, progress=gr.Progress(track_tqdm=True)):
114
  print("process start")
115
  if image_url:
116
  print(image_url)
 
120
 
121
  size = (orginal_image.size[0], orginal_image.size[1])
122
  print("image size", size)
123
+ depth_image = get_depth_map(orginal_image, progress)
124
  generator = torch.Generator().manual_seed(seed)
125
+ print(prompt, n_prompt, guidance_scale, num_steps, control_strength)
126
  generated_image = pipe(
127
  prompt=prompt,
128
  negative_prompt=n_prompt,
 
136
  ).images[0]
137
 
138
  if upload_to_s3:
139
+ url = upload_image_to_s3(generated_image, account_id, access_key, secret_key, bucket, progress)
140
  result = {"status": "success", "url": url}
141
  else:
142
  result = {"status": "success", "message": "Image generated but not uploaded"}
143
 
144
+ return [orginal_image, generated_image], json.dumps(result)
145
 
146
  with gr.Blocks() as demo:
147