Update app.py

app.py

@@ -178,7 +178,7 @@ def add_text(state, imagebox, textbox, image_process_mode):
 
     yield (state, state.to_gradio_chatbot(), "", None) + (disable_btn, disable_btn, disable_btn, enable_btn, enable_btn)
 
-def get_gen_images(out):
+def get_gen_images(out, pipe):
     pipe = pipe.to("cuda")
     img_embeds = out.image_embs
     if len(img_embeds) == 0:
@@ -211,7 +211,7 @@ def get_depth_images(out, org_size):
     grid_image = make_grid(depths, depth_layer_indices)
     return grid_image
 
-def get_seg_images(out, image):
+def get_seg_images(out, image, oneformer):
     oneformer = oneformer.to("cuda")
     seg_embs = out.seg_embs
     
@@ -284,10 +284,10 @@ def get_interm_outs(state):
         )
     
     depth_outs = get_depth_images(interm_outs, image_sizes[0]) 
-    seg_outs =  get_seg_images(interm_outs, images[0])
-    gen_outs = get_gen_images(interm_outs)
+    seg_outs =  get_seg_images(interm_outs, images[0], oneformer)
+    gen_outs = get_gen_images(interm_outs, pipe)
 
-    return images[0], images[0], images[0]
+    return depth_outs, seg_outs, gen_outs
     
 
 @spaces.GPU