more review fixes

Author: akhilg-nv

tripy/examples/diffusion/example.py

@@ -47,7 +47,7 @@ def compile_model(model, inputs, engine_path, verbose=False):
     if verbose:
         compile_end_time = time.perf_counter()
         print(f"saved engine to {engine_path}.")
-        print(f"took {compile_end_time - compile_start_time} seconds.")
+        print(f"Took {compile_end_time - compile_start_time} seconds.")
 
     return compiled_model
 
@@ -133,9 +133,7 @@ def save_image(image, args):
 
     # Save image
     print(f"[I] Saving image to {filename}")
-    if not os.path.isdir(os.path.dirname(filename)):
-        print(f"[I] Creating '{os.path.dirname(filename)}' directory.")
-        os.makedirs(os.path.dirname(filename))
+    os.makedirs(os.path.dirname(filename), exist_ok=True)
     image.save(filename)
 
 

tripy/examples/diffusion/models/clip_model.py

@@ -19,7 +19,7 @@
 
 from dataclasses import dataclass
 
-from examples.diffusion.helper import scaled_dot_product_attention
+from examples.diffusion.models.utils import scaled_dot_product_attention
 
 
 @dataclass

tripy/examples/diffusion/models/model.py

@@ -26,12 +26,12 @@
 from examples.diffusion.models.clip_model import CLIPTextTransformer, CLIPConfig
 from examples.diffusion.models.unet_model import UNetModel, UNetConfig
 from examples.diffusion.models.vae_model import AutoencoderKL, VAEConfig
-from examples.diffusion.helper import clamp
+from examples.diffusion.models.utils import clamp
 
 
 @dataclass
 class StableDiffusionConfig:
-    dtype: tp.dtype = tp.float32
+    dtype: tp.dtype
     clip_config: Optional[CLIPConfig] = field(default=None, init=False)
     unet_config: Optional[UNetConfig] = field(default=None, init=False)
     vae_config: Optional[VAEConfig] = field(default=None, init=False)

tripy/examples/diffusion/models/unet_model.py

@@ -21,8 +21,7 @@
 import nvtripy as tp
 from dataclasses import dataclass
 
-from examples.diffusion.helper import scaled_dot_product_attention
-from examples.diffusion.models.vae_model import Upsample, Downsample
+from examples.diffusion.models.utils import scaled_dot_product_attention, Upsample, Downsample
 
 
 @dataclass

tripy/examples/diffusion/helper.py renamed to tripy/examples/diffusion/models/utils.py

@@ -38,3 +38,21 @@ def scaled_dot_product_attention(
 
 def clamp(tensor: tp.Tensor, min: int, max: int):
     return tp.minimum(tp.maximum(tensor, min), max)
+
+
+class Upsample(tp.Module):
+    def __init__(self, config, channels):
+        self.conv = tp.Conv(channels, channels, (3, 3), padding=((1, 1), (1, 1)), dtype=config.dtype)
+
+    def __call__(self, x):
+        bs, c, py, px = x.shape
+        x = tp.reshape(tp.expand(tp.reshape(x, (bs, c, py, 1, px, 1)), (bs, c, py, 2, px, 2)), (bs, c, py * 2, px * 2))
+        return self.conv(x)
+
+
+class Downsample(tp.Module):
+    def __init__(self, config, channels):
+        self.conv = tp.Conv(channels, channels, (3, 3), stride=(2, 2), padding=((1, 1), (1, 1)), dtype=config.dtype)
+
+    def __call__(self, x):
+        return self.conv(x)

tripy/examples/diffusion/models/vae_model.py

@@ -20,7 +20,7 @@
 import nvtripy as tp
 from dataclasses import dataclass
 
-from examples.diffusion.helper import scaled_dot_product_attention
+from examples.diffusion.models.utils import scaled_dot_product_attention, Upsample, Downsample
 
 
 @dataclass
@@ -80,24 +80,6 @@ def __call__(self, x):
         return self.conv_shortcut(x) + h
 
 
-class Downsample(tp.Module):
-    def __init__(self, config, channels):
-        self.conv = tp.Conv(channels, channels, (3, 3), stride=(2, 2), padding=((1, 1), (1, 1)), dtype=config.dtype)
-
-    def __call__(self, x):
-        return self.conv(x)
-
-
-class Upsample(tp.Module):
-    def __init__(self, config, channels):
-        self.conv = tp.Conv(channels, channels, (3, 3), padding=((1, 1), (1, 1)), dtype=config.dtype)
-
-    def __call__(self, x):
-        bs, c, py, px = x.shape
-        x = tp.reshape(tp.expand(tp.reshape(x, (bs, c, py, 1, px, 1)), (bs, c, py, 2, px, 2)), (bs, c, py * 2, px * 2))
-        return self.conv(x)
-
-
 class UpDecoderBlock2D(tp.Module):
     def __init__(self, config: VAEConfig, start_channels, channels, use_upsampler=True):
         self.resnets = [