add nparams count

app.py

@@ -50,20 +50,25 @@ if __name__ == "__main__":
                 with gr.Tab("Example Prompts"):
                     examples = gr.Examples(examples=example_list, inputs=[text])
 
-            with gr.Column(variant='panel', scale=35):
+            with gr.Column(variant='panel',scale=35):
                 # Define original model output components
                 gr.Markdown('<h2 align="center">Original Stable Diffusion 1.4</h2>')
                 original_model_output = gr.Image(label="Original Model")
                 with gr.Row().style(equal_height=True):
-                    original_model_test_time = gr.Textbox(value="", label="Inference Time (sec)")
+                    with gr.Column():
+                        original_model_test_time = gr.Textbox(value="", label="Inference Time (sec)")
+                        original_model_params = gr.Textbox(value=servicer.get_sdm_params(servicer.pipe_original), label="# Parameters")
                     original_model_error = gr.Markdown()
+                
 
-            with gr.Column(variant='panel', scale=35):
+            with gr.Column(variant='panel',scale=35):
                 # Define compressed model output components
                 gr.Markdown('<h2 align="center">Compressed Stable Diffusion (Ours)</h2>')
-                compressed_model_output = gr.Image(label="Compressed Model")
+                compressed_model_output = gr.Image(label="Compressed Model")                
                 with gr.Row().style(equal_height=True):
-                    compressed_model_test_time = gr.Textbox(value="", label="Inference Time (sec)")
+                    with gr.Column():
+                        compressed_model_test_time = gr.Textbox(value="", label="Inference Time (sec)")
+                        compressed_model_params = gr.Textbox(value=servicer.get_sdm_params(servicer.pipe_compressed), label="# Parameters")
                     compressed_model_error = gr.Markdown()
 
         inputs = [text, negative, guidance_scale, steps, seed]

demo.py

@@ -26,6 +26,17 @@ class SdmCompressionDemo:
             self.pipe_compressed = self.pipe_compressed.to(self.device)
         self.device_msg = 'Tested on GPU.' if 'cuda' in self.device else 'Tested on CPU.'
 
+    def _count_params(self, model):
+        return sum(p.numel() for p in model.parameters())
+
+    def get_sdm_params(self, pipe):
+        params_unet = self._count_params(pipe.unet)
+        params_text_enc = self._count_params(pipe.text_encoder)
+        params_image_dec = self._count_params(pipe.vae.decoder)
+        params_total = params_unet + params_text_enc + params_image_dec
+        return f"Total {(params_total/1e6):.1f}M (U-Net {(params_unet/1e6):.1f}M)"
+
+
     def generate_image(self, pipe, text, negative, guidance_scale, steps, seed):
         generator = torch.Generator(self.device).manual_seed(seed)
         start = time.time()