applying gr.State

app.py

@@ -239,7 +239,7 @@ class Experiment(Component):
         idx = [metric.__class__.__name__ for metric in metric_info[0]].index(metric_name)
         return metric_info[1][idx]
 
-    def generate_record(self, data_id, metric_names):
+    def generate_record(self, checkbox_group_info, data_id, metric_names):
         record = {}
         _base = self.experiment.run_batch([data_id], 0, 0, 0)
         record['data_id'] = data_id
@@ -252,7 +252,7 @@ class Experiment(Component):
         metrics_ids = [self.get_metric_id_by_name(metric_nm) for metric_nm in metric_names]
 
         cnt = 0
-        for info in self.explainer_checkbox_group.info:
+        for info in checkbox_group_info:
             if info['checked']:
                 base = self.experiment.run_batch([data_id], info['id'], info['pp_id'], 0)
                 record['explanations'].append({
@@ -334,9 +334,9 @@ class Experiment(Component):
                     plot = gr.Image(value=None, label="Blank", visible=False)
                     plots.append(plot)
 
-        def show_plots():
+        def show_plots(checkbox_group_info):
             _plots = [gr.Textbox(label="Prediction result", visible=False)]
-            num_plots = sum([1 for info in self.explainer_checkbox_group.info if info['checked']])
+            num_plots = sum([1 for info in checkbox_group_info if info['checked']])
             n_rows = num_plots // PLOT_PER_LINE
             n_rows = n_rows + 1 if num_plots % PLOT_PER_LINE != 0 else n_rows
             _plots += [gr.Image(value=None, label="Blank", visible=True)] * (n_rows * PLOT_PER_LINE)
@@ -344,7 +344,7 @@ class Experiment(Component):
             return _plots
         
         @spaces.GPU
-        def render_plots(data_id, *metric_inputs):
+        def render_plots(data_id, checkbox_group_info, *metric_inputs):
             # Clear Cache Files
             # print(f"GPU Check: {torch.cuda.is_available()}")
             # print("Which GPU: ", torch.cuda.current_device())
@@ -360,12 +360,15 @@ class Experiment(Component):
                 if metric:
                     metric_input += metric
                     
-            record = self.generate_record(data_id, metric_input)
+            record = self.generate_record(checkbox_group_info, data_id, metric_input)
 
             pred = self.get_prediction(record)
             plots = [gr.Textbox(label="Prediction result", value=pred, visible=True)]
 
-            num_plots = sum([1 for info in self.explainer_checkbox_group.info if info['checked']])
+            # for info in checkbox_group_info:
+            #     if info['checked']:
+            #         print(info)
+            num_plots = sum([1 for info in checkbox_group_info if info['checked']])
             n_rows = num_plots // PLOT_PER_LINE
             n_rows = n_rows + 1 if num_plots % PLOT_PER_LINE != 0 else n_rows
 
@@ -383,8 +386,8 @@ class Experiment(Component):
 
             return plots
         
-        bttn.click(show_plots, outputs=plots)
-        bttn.click(render_plots, inputs=[data_id] + metric_inputs, outputs=plots)
+        bttn.click(show_plots, inputs=[self.explainer_checkbox_group.info], outputs=plots)
+        bttn.click(render_plots, inputs=[data_id, self.explainer_checkbox_group.info] + metric_inputs, outputs=plots)
 
 
 
@@ -397,30 +400,33 @@ class ExplainerCheckboxGroup(Component):
         self.gallery = gallery
         explainers, exp_ids = self.experiment.manager.get_explainers()
 
-        self.info = []
+        info = []
         for exp, exp_id in zip(explainers, exp_ids):
             exp_nm = exp.__class__.__name__
             if exp_nm in DEFAULT_EXPLAINER:
                 checked = True
             else:
                 checked = False
-            self.info.append({'nm': exp_nm, 'id': exp_id, 'pp_id' : 0, 'mode': 'default', 'checked': checked})
+            info.append({'nm': exp_nm, 'id': exp_id, 'pp_id' : 0, 'mode': 'default', 'checked': checked})
+        self.static_info = sorted(info, key=lambda x: (x['nm'] not in DEFAULT_EXPLAINER, x['nm']))
+        self.info = gr.State(info)
 
-    def update_check(self, exp_id, val=None):
-        for info in self.info:
+    def update_check(self, checkbox_group_info, exp_id, val=None):
+        for info in checkbox_group_info:
             if info['id'] == exp_id:
                 if val is not None:
                     info['checked'] = val
                 else:
                     info['checked'] = not info['checked']
+        return checkbox_group_info
 
-    def insert_check(self, exp_nm, exp_id, pp_id):
-        if exp_id in [info['id'] for info in self.info]:
+    def insert_check(self, checkbox_group_info, exp_nm, exp_id, pp_id):
+        if exp_id in [info['id'] for info in checkbox_group_info]:
             return
+        checkbox_group_info.append({'nm': exp_nm, 'id': exp_id, 'pp_id' : pp_id, 'mode': 'optimal', 'checked': False})
+        return checkbox_group_info
 
-        self.info.append({'nm': exp_nm, 'id': exp_id, 'pp_id' : pp_id, 'mode': 'optimal', 'checked': False})
-
-    def update_gallery_change(self):
+    def update_gallery_change(self, checkbox_group_info):
         checkboxes = []
         bttns = []
         for exp in self.explainer_objs:
@@ -431,10 +437,10 @@ class ExplainerCheckboxGroup(Component):
 
         for exp in self.explainer_objs:
             val = exp.explainer_name in DEFAULT_EXPLAINER
-            self.update_check(exp.default_exp_id, val)
+            checkbox_group_info = self.update_check(checkbox_group_info, exp.default_exp_id, val)
             if hasattr(exp, "optimal_exp_id"):
-                self.update_check(exp.optimal_exp_id, False)
-        return checkboxes + bttns
+                checkbox_group_info = self.update_check(checkbox_group_info, exp.optimal_exp_id, False)
+        return checkboxes + bttns + [checkbox_group_info]
 
     def get_checkboxes(self):
         checkboxes = []
@@ -447,11 +453,10 @@ class ExplainerCheckboxGroup(Component):
     
     def show(self):
         cnt = 0
-        sorted_info = sorted(self.info, key=lambda x: (x['nm'] not in DEFAULT_EXPLAINER, x['nm']))
         with gr.Accordion("Explainers", open=True):
             while cnt * PLOT_PER_LINE < len(self.explainer_names):
                 with gr.Row():
-                    for info in sorted_info[cnt*PLOT_PER_LINE:(cnt+1)*PLOT_PER_LINE]:
+                    for info in self.static_info[cnt*PLOT_PER_LINE:(cnt+1)*PLOT_PER_LINE]:
                         explainer_obj = ExplainerCheckbox(info['nm'], self, self.experiment, self.gallery)
                         self.explainer_objs.append(explainer_obj)
                         explainer_obj.show()
@@ -461,7 +466,8 @@ class ExplainerCheckboxGroup(Component):
         bttns = self.get_bttns()
         self.gallery.gallery_obj.select(
             fn=self.update_gallery_change,
-            outputs=checkboxes + bttns
+            inputs=self.info,
+            outputs=checkboxes + bttns + [self.info],
         )
 
     
@@ -488,28 +494,31 @@ class ExplainerCheckbox(Component):
     def get_str_ppid(self, pp_obj):
         return pp_obj.pooling_fn.__class__.__name__ + pp_obj.normalization_fn.__class__.__name__
 
-    def default_on_select(self, evt: gr.EventData):
-        self.groups.update_check(self.default_exp_id, evt._data['value'])
+    def default_on_select(self, evt: gr.EventData, checkbox_group_info):
+        checkbox_group_info = self.groups.update_check(checkbox_group_info, self.default_exp_id, evt._data['value'])
+        return checkbox_group_info
 
-    def optimal_on_select(self, evt: gr.EventData):
+    def optimal_on_select(self, evt: gr.EventData, checkbox_group_info):
         if hasattr(self, "optimal_exp_id"):
-            self.groups.update_check(self.optimal_exp_id, evt._data['value'])
+            checkbox_group_info = self.groups.update_check(checkbox_group_info, self.optimal_exp_id, evt._data['value'])
         else:
             raise ValueError("Optimal explainer id is not found.")
+        return checkbox_group_info
 
     def show(self):
         val = self.explainer_name in DEFAULT_EXPLAINER
         with gr.Accordion(self.explainer_name, open=val):
-            checked = next(filter(lambda x: x['nm'] == self.explainer_name, self.groups.info))['checked']
+            checked = next(filter(lambda x: x['nm'] == self.explainer_name, self.groups.static_info))['checked']
             self.default_check = gr.Checkbox(label="Default Parameter", value=checked, interactive=True)
             self.opt_check = gr.Checkbox(label="Optimized Parameter (Not Optimal)", interactive=False)
 
-            self.default_check.select(self.default_on_select)
-            self.opt_check.select(self.optimal_on_select)
+            self.default_check.select(self.default_on_select, self.groups.info, self.groups.info)
+            self.opt_check.select(self.optimal_on_select, self.groups.info, self.groups.info)
 
             self.bttn = gr.Button(value="Optimize", size="sm", variant="primary")
 
-            def optimize():
+            @spaces.GPU
+            def optimize(checkbox_group_info):
                 data_id = self.gallery.selected_index
         
                 opt_output = self.experiment.optimize(
@@ -527,18 +536,18 @@ class ExplainerCheckbox(Component):
                         opt_postprocessor_id = pp_id
                         break
 
-                opt_explainer_id = max([x['id'] for x in self.groups.info]) + 1
+                opt_explainer_id = max([x['id'] for x in checkbox_group_info]) + 1
                 opt_output.explainer.model = self.experiment.model
                 self.experiment.manager._explainers.append(opt_output.explainer)
                 self.experiment.manager._explainer_ids.append(opt_explainer_id)
-                self.groups.insert_check(self.explainer_name, opt_explainer_id, opt_postprocessor_id)
+                self.groups.insert_check(checkbox_group_info, self.explainer_name, opt_explainer_id, opt_postprocessor_id)
                 self.optimal_exp_id = opt_explainer_id
                 checkbox = gr.update(label="Optimized Parameter (Optimal)", interactive=True)
                 bttn = gr.update(value="Optimized", variant="secondary")
 
-                return [checkbox, bttn]
+                return [checkbox_group_info, checkbox, bttn]
 
-            self.bttn.click(optimize, outputs=[self.opt_check, self.bttn], queue=True, concurrency_limit=1)
+            self.bttn.click(optimize, inputs=[self.groups.info], outputs=[self.groups.info, self.opt_check, self.bttn], queue=True, concurrency_limit=1)
         
 
 class ExpRes(Component):
@@ -692,3 +701,4 @@ experiments['experiment2'] = {
 app = ImageClsApp(experiments)
 demo = app.launch()
 demo.launch(favicon_path=f"static/XAI-Top-PnP.svg")
+# demo.launch(favicon_path=f"static/XAI-Top-PnP.svg", share=True)