Changed device to global variable in gradio (ie gr.State instance)

app.py

@@ -208,26 +208,26 @@ def get_ind_to_filter(text, word_ids, keywords):
     return inds_to_filter
 
 @spaces.GPU
-def count(image, text, prompts, state):
+def count(image, text, prompts, state, device):
     print("state: " + str(state))
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None:
         prompts = {"image": image, "points": []}
     input_image, _ = transform(image, {"exemplars": torch.tensor([])})
-    input_image = input_image.unsqueeze(0).to(args.device)
+    input_image = input_image.unsqueeze(0).to(device)
     exemplars = get_box_inputs(prompts["points"])
     print(exemplars)
     input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
-    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(args.device)
-    exemplars = [exemplars["exemplars"].to(args.device)]
+    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
+    exemplars = [exemplars["exemplars"].to(device)]
     
     with torch.no_grad():
         model_output = model(
                 nested_tensor_from_tensor_list(input_image),
                 nested_tensor_from_tensor_list(input_image_exemplars),
                 exemplars,
-                [torch.tensor([0]).to(args.device) for _ in range(len(input_image))],
+                [torch.tensor([0]).to(device) for _ in range(len(input_image))],
                 captions=[text + " ."] * len(input_image),
             )
     
@@ -297,25 +297,25 @@ def count(image, text, prompts, state):
     return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]), new_submit_btn, gr.Tab(visible=True), step_3, state)
 
 @spaces.GPU
-def count_main(image, text, prompts):
+def count_main(image, text, prompts, device):
     keywords = "" # do not handle this for now
     # Handle no prompt case.
     if prompts is None:
         prompts = {"image": image, "points": []}
     input_image, _ = transform(image, {"exemplars": torch.tensor([])})
-    input_image = input_image.unsqueeze(0).to(args.device)
+    input_image = input_image.unsqueeze(0).to(device)
     exemplars = get_box_inputs(prompts["points"])
     print(exemplars)
     input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
-    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(args.device)
-    exemplars = [exemplars["exemplars"].to(args.device)]
+    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
+    exemplars = [exemplars["exemplars"].to(device)]
     
     with torch.no_grad():
         model_output = model(
                 nested_tensor_from_tensor_list(input_image),
                 nested_tensor_from_tensor_list(input_image_exemplars),
                 exemplars,
-                [torch.tensor([0]).to(args.device) for _ in range(len(input_image))],
+                [torch.tensor([0]).to(device) for _ in range(len(input_image))],
                 captions=[text + " ."] * len(input_image),
             )
     
@@ -396,6 +396,7 @@ As shown earlier, there are 3 ways to specify the object to count: (1) with text
 
 with gr.Blocks(title="CountGD: Multi-Modal Open-World Counting", theme="soft", head="""<meta name="viewport" content="width=device-width, initial-scale=1, user-scalable=1">""") as demo:
     state = gr.State(value=[AppSteps.JUST_TEXT])
+    device = gr.State(args.device)
     with gr.Tab("Tutorial"):
         with gr.Row():
             with gr.Column():
@@ -419,7 +420,7 @@ with gr.Blocks(title="CountGD: Multi-Modal Open-World Counting", theme="soft", h
             pred_count = gr.Number(label="Predicted Count", visible=False)
         submit_btn = gr.Button("Count", variant="primary", interactive=True)
                 
-        submit_btn.click(fn=remove_label, inputs=[detected_instances], outputs=[detected_instances]).then(fn=count, inputs=[input_image, input_text, exemplar_image, state], outputs=[detected_instances, pred_count, submit_btn, step_2, step_3, state])
+        submit_btn.click(fn=remove_label, inputs=[detected_instances], outputs=[detected_instances]).then(fn=count, inputs=[input_image, input_text, exemplar_image, state, device], outputs=[detected_instances, pred_count, submit_btn, step_2, step_3, state])
         exemplar_image.change(check_submit_btn, inputs=[exemplar_image, state], outputs=[submit_btn])
     with gr.Tab("App", visible=True) as main_app:
         
@@ -445,7 +446,7 @@ with gr.Blocks(title="CountGD: Multi-Modal Open-World Counting", theme="soft", h
               submit_btn_main = gr.Button("Count", variant="primary")
               clear_btn_main = gr.ClearButton(variant="secondary")
         gr.Examples(label="Examples: click on a row to load the example. Add visual exemplars by drawing boxes on the loaded \"Visual Exemplar Image.\"", examples=examples, inputs=[input_image_main, input_text_main, exemplar_image_main])
-        submit_btn_main.click(fn=remove_label, inputs=[detected_instances_main], outputs=[detected_instances_main]).then(fn=count_main, inputs=[input_image_main, input_text_main, exemplar_image_main], outputs=[detected_instances_main, pred_count_main])
+        submit_btn_main.click(fn=remove_label, inputs=[detected_instances_main], outputs=[detected_instances_main]).then(fn=count_main, inputs=[input_image_main, input_text_main, exemplar_image_main, device], outputs=[detected_instances_main, pred_count_main])
         clear_btn_main.add([input_image_main, input_text_main, exemplar_image_main, detected_instances_main, pred_count_main])