Update model

models.py

@@ -1,14 +1,3 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-# --------------------------------------------------------
-# References:
-# GLIDE: https://github.com/openai/glide-text2im
-# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py
-# --------------------------------------------------------
-
 import torch
 import torch.nn as nn
 import numpy as np
@@ -16,6 +5,8 @@ import math
 from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
 import einops
 
+from huggingface_hub import PyTorchModelHubMixin
+
 import torch.utils.checkpoint as checkpoint
 
 from transformers import PreTrainedModel
@@ -371,7 +362,7 @@ class FinalLayer(nn.Module):
         return x
 
 
-class FLAV(nn.Module):
+class FLAV(nn.Module, PyTorchModelHubMixin):
     """
     Diffusion model with a Transformer backbone.
     """
@@ -748,4 +739,3 @@ FLAV_models = {
     'FLAV-B/1' :  FLAV_B_1,   'FLAV-B/2':  FLAV_B_2,   'FLAV-B/4':  FLAV_B_4,   'FLAV-B/8':  FLAV_B_8,
     'FLAV-S/2' :  FLAV_S_2,   'FLAV-S/4':  FLAV_S_4,   'FLAV-S/8':  FLAV_S_8,
 }
-