Update app.py

app.py

@@ -66,8 +66,9 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return seed
 
 
-def get_depth_map(image, progress):
+def get_depth_map(image):
     original_size = (image.size[1], image.size[0]) 
+    print("start generate depth", original_size)
     image = feature_extractor(images=image, return_tensors="pt").pixel_values.to("cuda")
     with torch.no_grad(), torch.autocast("cuda"):
         depth_map = depth_estimator(image).predicted_depth
@@ -83,10 +84,11 @@ def get_depth_map(image, progress):
     image = torch.cat([depth_map] * 3, dim=1)
     image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
     image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
+    print("generate depth success")
     return image
 
 
-def upload_image_to_s3(image, account_id, access_key, secret_key, bucket_name, progress):
+def upload_image_to_s3(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"
 
@@ -119,10 +121,11 @@ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, contr
         orginal_image = Image.fromarray(image)
 
     size = (orginal_image.size[0], orginal_image.size[1])
-    print("image size", size)
+    print("gorinal image size", size)
     depth_image = get_depth_map(orginal_image, progress)
     generator = torch.Generator().manual_seed(seed)
     print(prompt, n_prompt, guidance_scale, num_steps, control_strength)
+    print("run pipe")
     generated_image = pipe(
         prompt=prompt,
         negative_prompt=n_prompt,
@@ -134,14 +137,14 @@ def process(image, image_url, prompt, n_prompt, num_steps, guidance_scale, contr
         generator=generator,
         image=depth_image
     ).images[0]
-
+    print("geneate image success")
     if upload_to_s3:
         url = upload_image_to_s3(generated_image, account_id, access_key, secret_key, bucket, progress)
         result = {"status": "success", "url": url}
     else:
         result = {"status": "success", "message": "Image generated but not uploaded"}
     
-    return [orginal_image, generated_image], json.dumps(result)
+    return generated_image, json.dumps(result)
 
 with gr.Blocks() as demo:
     
@@ -172,7 +175,7 @@ with gr.Blocks() as demo:
         
 
         with gr.Column():
-            images = ImageSlider(label="Generate images", type="pil", slider_color="pink")
+            result = gr.Image(label="Generated Image")
             logs = gr.Textbox(label="logs")
             
     inputs = [
@@ -199,8 +202,8 @@ with gr.Blocks() as demo:
         ).then(
             fn=process,
             inputs=inputs,
-            outputs=[images, logs],
-            api_name=False
+            outputs=[result, logs],
+            api_name="predict"
         )
 
 demo.queue().launch()