Fix type annotation mismatches in pil_utils.py and outputs.py

Signed-off-by: Kunjan patel <[email protected]>

Author: coolkp

src/maxdiffusion/utils/outputs.py

@@ -16,7 +16,7 @@
 """
 
 from collections import OrderedDict
-from dataclasses import fields, is_dataclass
+from dataclasses import dataclass, fields, is_dataclass
 from typing import Any, Tuple
 
 import numpy as np
@@ -37,6 +37,7 @@ def is_tensor(x):
   return isinstance(x, np.ndarray)
 
 
+@dataclass
 class BaseOutput(OrderedDict):
   """
   Base class for all model outputs as dataclass. Has a `__getitem__` that allows indexing by integer or slice (like a

src/maxdiffusion/utils/pil_utils.py

@@ -16,11 +16,11 @@
   }
 else:
   PIL_INTERPOLATION = {
-      "linear": PIL.Image.LINEAR,
-      "bilinear": PIL.Image.BILINEAR,
-      "bicubic": PIL.Image.BICUBIC,
-      "lanczos": PIL.Image.LANCZOS,
-      "nearest": PIL.Image.NEAREST,
+      "linear": PIL.Image.LINEAR,  # pytype: disable=module-attr
+      "bilinear": PIL.Image.BILINEAR,  # pytype: disable=module-attr
+      "bicubic": PIL.Image.BICUBIC,  # pytype: disable=module-attr
+      "lanczos": PIL.Image.LANCZOS,  # pytype: disable=module-attr
+      "nearest": PIL.Image.NEAREST,  # pytype: disable=module-attr
   }
 
 
@@ -50,7 +50,7 @@ def numpy_to_pil(images):
   return pil_images
 
 
-def make_image_grid(images: List[PIL.Image.Image], rows: int, cols: int, resize: int = None) -> PIL.Image.Image:
+def make_image_grid(images: List[PIL.Image.Image], rows: int, cols: int, resize: int | None = None) -> PIL.Image.Image:
   """
   Prepares a single grid of images. Useful for visualization purposes.
   """