torch_dtype

app.py

@@ -327,7 +327,8 @@ with gr.Blocks(css=css) as demo:
             fn=generate_image,
             inputs=[prompt],
             outputs=[result, seed],
-            cache_examples="lazy",
+            cache_examples=True,
+            cache_mode="lazy",
         )
 
         gr.on(

tim/models/nvidia_radio/radio/extra_models.py

@@ -13,7 +13,7 @@ from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
 from .forward_intermediates import forward_intermediates
 from .input_conditioner import InputConditioner
 
-_has_torch_sdpa = hasattr(F, 'scaled_dot_product_attention')
+_has_torch_sdpa = hasattr(F, "scaled_dot_product_attention")
 
 
 class PaliGemmaWrapper(nn.Module):
@@ -52,18 +52,25 @@ class PaliGemmaWrapper(nn.Module):
         return self(x)
 
 
-def _get_paligemma_model(repo: str, embed_dim: int = None, dtype: torch.dtype = torch.bfloat16):
-    from transformers import PaliGemmaForConditionalGeneration, __version__ as tx_version
+def _get_paligemma_model(
+    repo: str, embed_dim: int = None, dtype: torch.dtype = torch.bfloat16
+):
+    from transformers import (
+        PaliGemmaForConditionalGeneration,
+        __version__ as tx_version,
+    )
 
-    if LooseVersion(tx_version) > LooseVersion('4.44.2'):
-        warnings.warn(f'Your transformers version "{tx_version}" is higher than 4.44.2, and for whatever reason, PaliGemma might be broken.')
+    if LooseVersion(tx_version) > LooseVersion("4.44.2"):
+        warnings.warn(
+            f'Your transformers version "{tx_version}" is higher than 4.44.2, and for whatever reason, PaliGemma might be broken.'
+        )
 
     extra_args = dict()
 
     if dtype is not None:
-        extra_args['torch_dtype'] = dtype
-        rev = str(dtype).split('.')[-1]
-        extra_args['revision'] = rev
+        extra_args["dtype"] = dtype
+        rev = str(dtype).split(".")[-1]
+        extra_args["revision"] = rev
 
     model = PaliGemmaForConditionalGeneration.from_pretrained(repo, **extra_args)
 
@@ -73,22 +80,31 @@ def _get_paligemma_model(repo: str, embed_dim: int = None, dtype: torch.dtype =
 
     return vis_model
 
+
 @register_model
 def paligemma_896_student(**kwargs):
-    model = _get_paligemma_model('google/paligemma-3b-pt-896', embed_dim=1152, dtype=None)
+    model = _get_paligemma_model(
+        "google/paligemma-3b-pt-896", embed_dim=1152, dtype=None
+    )
 
     return model
 
 
 def dv2_sdpa(self, x: torch.Tensor) -> torch.Tensor:
     B, N, C = x.shape
-    qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+    qkv = (
+        self.qkv(x)
+        .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+        .permute(2, 0, 3, 1, 4)
+    )
 
     q, k, v = qkv[0], qkv[1], qkv[2]
     x = F.scaled_dot_product_attention(
-        q, k, v,
+        q,
+        k,
+        v,
         is_causal=False,
-        dropout_p=self.attn_drop.p if self.training else 0.,
+        dropout_p=self.attn_drop.p if self.training else 0.0,
         scale=self.scale,
     )
     x = x.transpose(1, 2).reshape(B, N, C)
@@ -96,11 +112,14 @@ def dv2_sdpa(self, x: torch.Tensor) -> torch.Tensor:
     x = self.proj_drop(x)
     return x
 
-def _load_dino_v2(dino_v2_model, cache_dir: Optional[str] = None, pretrained=True, **kwargs):
+
+def _load_dino_v2(
+    dino_v2_model, cache_dir: Optional[str] = None, pretrained=True, **kwargs
+):
     if cache_dir:
         torch.hub.set_dir(cache_dir)
     model: nn.Module = torch.hub.load(
-        'facebookresearch/dinov2',
+        "facebookresearch/dinov2",
         dino_v2_model,
         pretrained=pretrained,
         # **kwargs,
@@ -108,11 +127,12 @@ def _load_dino_v2(dino_v2_model, cache_dir: Optional[str] = None, pretrained=Tru
 
     if _has_torch_sdpa:
         for n, m in model.named_modules():
-            if n.endswith('.attn'):
+            if n.endswith(".attn"):
                 m.forward = MethodType(dv2_sdpa, m)
 
     return model
 
+
 class DinoWrapper(nn.Module):
     def __init__(self, dino_model: nn.Module):
         super().__init__()
@@ -130,11 +150,11 @@ class DinoWrapper(nn.Module):
 
     @property
     def num_cls_tokens(self):
-        return getattr(self.inner, 'num_tokens', 1)
+        return getattr(self.inner, "num_tokens", 1)
 
     @property
     def num_registers(self):
-        return getattr(self.inner, 'num_register_tokens', 0)
+        return getattr(self.inner, "num_register_tokens", 0)
 
     @property
     def num_summary_tokens(self):
@@ -147,8 +167,8 @@ class DinoWrapper(nn.Module):
     def forward(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
         parts = self.inner.forward_features(*args, **kwargs)
 
-        cls_token = parts['x_norm_clstoken']
-        features = parts['x_norm_patchtokens']
+        cls_token = parts["x_norm_clstoken"]
+        features = parts["x_norm_patchtokens"]
 
         return cls_token, features
 
@@ -157,12 +177,13 @@ class DinoWrapper(nn.Module):
         x = self.inner.blocks(x)
         x_norm = self.inner.norm(x)
 
-        return x_norm[:, 0], x_norm[:, self.num_summary_tokens:]
+        return x_norm[:, 0], x_norm[:, self.num_summary_tokens :]
 
     def patchify(self, x: torch.Tensor) -> torch.Tensor:
         return self.inner.prepare_tokens_with_masks(x)
 
-    def forward_intermediates(self,
+    def forward_intermediates(
+        self,
         x: torch.Tensor,
         norm: bool = False,
         **kwargs,
@@ -199,8 +220,9 @@ def _dino_student(arch: str, **kwargs):
 
 @register_model
 def dino_v2_l_student(**kwargs):
-    return _dino_student('dinov2_vitl14_reg', **kwargs)
+    return _dino_student("dinov2_vitl14_reg", **kwargs)
+
 
 @register_model
 def dino_v2_g_student(**kwargs):
-    return _dino_student('dinov2_vitg14_reg', **kwargs)
+    return _dino_student("dinov2_vitg14_reg", **kwargs)

tim/models/utils/text_encoders.py

@@ -13,14 +13,14 @@ def load_text_encoder(text_encoder_dir, device, weight_dtype):
             text_encoder_dir,
             attn_implementation="flash_attention_2",
             device_map="cpu",
-            torch_dtype=weight_dtype,
+            dtype=weight_dtype,
         ).model
     elif "t5" in text_encoder_dir:
         text_encoder = T5EncoderModel.from_pretrained(
             text_encoder_dir,
             attn_implementation="sdpa",
             device_map="cpu",
-            torch_dtype=weight_dtype,
+            dtype=weight_dtype,
         )
     else:
         raise NotImplementedError