add license

Author: lucidtronix

ml4h/models/train_diffusion.py

@@ -261,30 +261,30 @@ def train_diffusion_control_model(args, supervised=False):
     model.normalizer.adapt(feature_batch)
     if args.inspect_model:
         model.network.summary(print_fn=logging.info, expand_nested=True)
-        tf.keras.utils.plot_model(
-            model.network,
-            to_file=f"{args.output_folder}/{args.id}/architecture_{args.id}_unet.png",
-            show_shapes=True,
-            show_dtype=False,
-            show_layer_names=True,
-            rankdir="TB",
-            expand_nested=True,
-            dpi=args.dpi,
-            layer_range=None,
-            show_layer_activations=False,
-        )
-        tf.keras.utils.plot_model(
-            model.control_embed_model,
-            to_file=f"{args.output_folder}/{args.id}/architecture_{args.id}_control_embed.png",
-            show_shapes=True,
-            show_dtype=False,
-            show_layer_names=True,
-            rankdir="TB",
-            expand_nested=True,
-            dpi=args.dpi,
-            layer_range=None,
-            show_layer_activations=True,
-        )
+        # tf.keras.utils.plot_model(
+        #     model.network,
+        #     to_file=f"{args.output_folder}/{args.id}/architecture_{args.id}_unet.png",
+        #     show_shapes=True,
+        #     show_dtype=False,
+        #     show_layer_names=True,
+        #     rankdir="TB",
+        #     expand_nested=True,
+        #     dpi=args.dpi,
+        #     layer_range=None,
+        #     show_layer_activations=False,
+        # )
+        # tf.keras.utils.plot_model(
+        #     model.control_embed_model,
+        #     to_file=f"{args.output_folder}/{args.id}/architecture_{args.id}_control_embed.png",
+        #     show_shapes=True,
+        #     show_dtype=False,
+        #     show_layer_names=True,
+        #     rankdir="TB",
+        #     expand_nested=True,
+        #     dpi=args.dpi,
+        #     layer_range=None,
+        #     show_layer_activations=True,
+        # )
         prefix_value = f'{args.output_folder}{args.id}/learning_generations/'
         # Create a partial function with reseed and prefix pre-filled
         if model.input_map.axes() == 2: