Update gen_demo.py

gen_demo.py

@@ -103,7 +103,7 @@ gpt_model = GPT_models[params.gpt_model.target](**params.gpt_model.params).to(de
 load_checkpoint(ckpt_path, gpt_model)
 gpt_model.eval();
 
-@spaces.GPU(duration=120)
+# @spaces.GPU(duration=120)
 def viz_diff_slots(model, slots, nums, cfg=1.0, return_figs=False):
     n_slots_inf = []
     for num_slots_to_inference in nums:
@@ -120,7 +120,7 @@ num_slots_to_gen = 32
 
 # Function to generate image from class
 @spaces.GPU(duration=120)
-def generate_from_class(class_id, cfg_scale):
+def generate_from_class_and_viz_diff_slots(class_id, cfg_scale, nums, cfg=1.0, return_figs=False):
     with torch.no_grad():
         dtype = torch.float
         num_slots_to_gen = 32
@@ -136,7 +136,18 @@ def generate_from_class(class_id, cfg_scale):
             null_slots = ae_model.dit.null_cond.expand(slots_gen.shape[0], -1, -1)
             null_slots = null_slots[:, num_slots_to_gen:, :]
             slots_gen = torch.cat([slots_gen, null_slots], dim=1)
-    return slots_gen
+
+    slots = slots_gen
+    n_slots_inf = []
+    for num_slots_to_inference in nums:
+        drop_mask = ae_model.nested_sampler(slots.shape[0], device, num_slots_to_inference)
+        recon_n = ae_model.sample(slots, drop_mask=drop_mask, cfg=cfg)
+        n_slots_inf.append(recon_n)
+    drop_mask = ae_model.nested_sampler(slots.shape[0], device, 32)
+    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])
 
 with gr.Blocks() as demo:
     with gr.Row():
@@ -213,13 +224,15 @@ with gr.Blocks() as demo:
         while len(labels) < 3:
             labels.append(256)
         
-        # Generate the image based on the selected class
-        slots_gen = generate_from_class(selected_class_id, cfg_scale=slider_value)
+        # # Generate the image based on the selected class
+        # slots_gen = generate_from_class(selected_class_id, cfg_scale=slider_value)
         
-        recon = viz_diff_slots(ae_model, slots_gen, [32], cfg=slider_value)[0]
+        # recon = viz_diff_slots(ae_model, slots_gen, [32], cfg=slider_value)[0]
         
-        # Always generate the model decomposition for potential gallery display
-        model_decompose = viz_diff_slots(ae_model, slots_gen, labels, cfg=slider_value)
+        # # 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)
         
         if not show_gallery_value:
             # If only the image should be shown, return just the processed image