Update gen_demo.py

gen_demo.py

@@ -148,7 +148,7 @@ def generate_from_class_and_viz_diff_slots(class_id, cfg_scale, nums, cfg=1.0, r
     recon_n = ae_model.sample(slots, drop_mask=drop_mask, cfg=cfg)
     recon = recon_n
     
-    return slots_gen, [convert_np(n_slots_inf[i][0]) for i in range(len(n_slots_inf))], convert_np(recon[0])
+    return [convert_np(n_slots_inf[i][0]) for i in range(len(n_slots_inf))], convert_np(recon[0])
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -232,8 +232,12 @@ with gr.Blocks() as demo:
         
         # # Always generate the model decomposition for potential gallery display
         # model_decompose = viz_diff_slots(ae_model, slots_gen, labels, cfg=slider_value)
-        slots_gen, model_decompose, recon = \
-        generate_from_class_and_viz_diff_slots(selected_class_id, cfg_scale=slider_value, nums=labels, cfg=slider_value)
+        model_decompose, recon = generate_from_class_and_viz_diff_slots(
+            selected_class_id, 
+            cfg_scale=slider_value, 
+            nums=labels, 
+            cfg=slider_value
+        )
         
         if not show_gallery_value:
             # If only the image should be shown, return just the processed image