Update model.py

model.py

@@ -1,148 +1,155 @@
-import json
-import torch
-import torch.nn as nn
-import os
-from pathlib import Path
-from typing import Optional, Union, Dict
-from huggingface_hub import snapshot_download
-import warnings
-
-class ConvVAE(nn.Module):
- def __init__(self, latent_size):
- super(ConvVAE, self).__init__()
-
- # Encoder
- self.encoder = nn.Sequential(
- nn.Conv2d(3, 64, 3, stride=2, padding=1), # (batch, 64, 64, 64)
- nn.BatchNorm2d(64),
- nn.ReLU(),
- nn.Conv2d(64, 128, 3, stride=2, padding=1), # (batch, 128, 32, 32)
- nn.BatchNorm2d(128),
- nn.ReLU(),
- nn.Conv2d(128, 256, 3, stride=2, padding=1), # (batch, 256, 16, 16)
- nn.BatchNorm2d(256),
- nn.ReLU(),
- nn.Conv2d(256, 512, 3, stride=2, padding=1), # (batch, 512, 8, 8)
- nn.BatchNorm2d(512),
- nn.ReLU()
- )
-
- self.fc_mu = nn.Linear(512 * 8 * 8, latent_size)
- self.fc_logvar = nn.Linear(512 * 8 * 8, latent_size)
-
- self.fc2 = nn.Linear(latent_size, 512 * 8 * 8)
-
- self.decoder = nn.Sequential(
- nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1), # (batch, 256, 16, 16)
- nn.BatchNorm2d(256),
- nn.ReLU(),
- nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1), # (batch, 128, 32, 32)
- nn.BatchNorm2d(128),
- nn.ReLU(),
- nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1), # (batch, 64, 64, 64)
- nn.BatchNorm2d(64),
- nn.ReLU(),
- nn.ConvTranspose2d(64, 3, 4, stride=2, padding=1), # (batch, 3, 128, 128)
- nn.Tanh() 
- )
-
- def forward(self, x):
- mu, logvar = self.encode(x)
- z = self.reparameterize(mu, logvar)
- decoded = self.decode(z)
- return decoded, mu, logvar
-
- def encode(self, x):
- x = self.encoder(x)
- x = x.view(x.size(0), -1) 
- mu = self.fc_mu(x)
- logvar = self.fc_logvar(x)
- return mu, logvar
-
- def reparameterize(self, mu, logvar):
- std = torch.exp(0.5 * logvar)
- eps = torch.randn_like(std)
- return mu + eps * std
-
- def decode(self, z):
- x = self.fc2(z)
- x = x.view(-1, 512, 8, 8)
- decoded = self.decoder(x)
- return decoded
-
- @classmethod
- def from_pretrained(
- cls,
- model_id: str,
- revision: Optional[str] = None,
- cache_dir: Optional[Union[str, Path]] = None,
- force_download: bool = False,
- proxies: Optional[Dict] = None,
- resume_download: bool = False,
- local_files_only: bool = False,
- token: Union[str, bool, None] = None,
- map_location: str = "cpu",
- strict: bool = False,
- **model_kwargs,
- ):
- """
- Load a pretrained model from a given model ID.
-
- Args:
- model_id (str): Identifier of the model to load.
- revision (Optional[str]): Specific model revision to use.
- cache_dir (Optional[Union[str, Path]]): Directory to store downloaded models.
- force_download (bool): Force re-download even if the model exists.
- proxies (Optional[Dict]): Proxy configuration for downloads.
- resume_download (bool): Resume interrupted downloads.
- local_files_only (bool): Use only local files, don't download.
- token (Union[str, bool, None]): Token for API authentication.
- map_location (str): Device to map model to. Defaults to "cpu".
- strict (bool): Enforce strict state_dict loading.
- **model_kwargs: Additional keyword arguments for model initialization.
-
- Returns:
- An instance of the model loaded from the pretrained weights.
- """
- model_dir = Path(model_id)
- if not model_dir.exists():
- model_dir = Path(
- snapshot_download(
- repo_id=model_id,
- revision=revision,
- cache_dir=cache_dir,
- force_download=force_download,
- proxies=proxies,
- resume_download=resume_download,
- token=token,
- local_files_only=local_files_only,
- )
- )
-
- config_file = model_dir / "config.json"
- with open(config_file, 'r') as f:
- config = json.load(f)
-
- latent_size = config.get('latent_size')
- if latent_size is None:
- raise ValueError("The configuration file is missing the 'latent_size' key.")
-
- model = cls(latent_size, **model_kwargs)
-
- model_file = model_dir / "model_conv_vae_256_epoch_304.pth"
- if not model_file.exists():
- raise FileNotFoundError(f"The model checkpoint '{model_file}' does not exist.")
- 
- state_dict = torch.load(model_file, map_location=map_location)
-
- new_state_dict = {}
- for k, v in state_dict.items():
- if k.startswith('_orig_mod.'):
- new_state_dict[k[len('_orig_mod.'):]] = v
- else:
- new_state_dict[k] = v
-
- model.load_state_dict(new_state_dict, strict=strict)
- model.to(map_location)
-
- return model
+import json
+import torch
+import torch.nn as nn
+import os
+from pathlib import Path
+from typing import Optional, Union, Dict
+from huggingface_hub import snapshot_download
+import warnings
+
+class ConvVAE(nn.Module):
+ def __init__(self, latent_size):
+ super(ConvVAE, self).__init__()
+
+ # Encoder
+ self.encoder = nn.Sequential(
+ nn.Conv2d(3, 64, 3, stride=2, padding=1), # (batch, 64, 64, 64)
+ nn.BatchNorm2d(64),
+ nn.ReLU(),
+ nn.Conv2d(64, 128, 3, stride=2, padding=1), # (batch, 128, 32, 32)
+ nn.BatchNorm2d(128),
+ nn.ReLU(),
+ nn.Conv2d(128, 256, 3, stride=2, padding=1), # (batch, 256, 16, 16)
+ nn.BatchNorm2d(256),
+ nn.ReLU(),
+ nn.Conv2d(256, 512, 3, stride=2, padding=1), # (batch, 512, 8, 8)
+ nn.BatchNorm2d(512),
+ nn.ReLU()
+ )
+
+ self.fc_mu = nn.Linear(512 * 8 * 8, latent_size)
+ self.fc_logvar = nn.Linear(512 * 8 * 8, latent_size)
+
+ self.fc2 = nn.Linear(latent_size, 512 * 8 * 8)
+
+ self.decoder = nn.Sequential(
+ nn.ConvTranspose2d(512, 256, 4, stride=2, padding=1), # (batch, 256, 16, 16)
+ nn.BatchNorm2d(256),
+ nn.ReLU(),
+ nn.ConvTranspose2d(256, 128, 4, stride=2, padding=1), # (batch, 128, 32, 32)
+ nn.BatchNorm2d(128),
+ nn.ReLU(),
+ nn.ConvTranspose2d(128, 64, 4, stride=2, padding=1), # (batch, 64, 64, 64)
+ nn.BatchNorm2d(64),
+ nn.ReLU(),
+ nn.ConvTranspose2d(64, 3, 4, stride=2, padding=1), # (batch, 3, 128, 128)
+ nn.Tanh() 
+ )
+
+ def forward(self, x):
+ mu, logvar = self.encode(x)
+ z = self.reparameterize(mu, logvar)
+ decoded = self.decode(z)
+ return decoded, mu, logvar
+
+ def encode(self, x):
+ x = self.encoder(x)
+ x = x.view(x.size(0), -1) 
+ mu = self.fc_mu(x)
+ logvar = self.fc_logvar(x)
+ return mu, logvar
+
+ def reparameterize(self, mu, logvar):
+ std = torch.exp(0.5 * logvar)
+ eps = torch.randn_like(std)
+ return mu + eps * std
+
+ def decode(self, z):
+ x = self.fc2(z)
+ x = x.view(-1, 512, 8, 8)
+ decoded = self.decoder(x)
+ return decoded
+
+ @classmethod
+ def from_pretrained(
+ cls,
+ model_id: str,
+ revision: Optional[str] = None,
+ cache_dir: Optional[Union[str, Path]] = None,
+ force_download: bool = False,
+ proxies: Optional[Dict] = None,
+ resume_download: bool = False,
+ local_files_only: bool = False,
+ token: Union[str, bool, None] = None,
+ map_location: str = "cpu",
+ strict: bool = False,
+ **model_kwargs,
+ ):
+ """
+ Load a pretrained model from a given model ID.
+
+ Args:
+ model_id (str): Identifier of the model to load.
+ revision (Optional[str]): Specific model revision to use.
+ cache_dir (Optional[Union[str, Path]]): Directory to store downloaded models.
+ force_download (bool): Force re-download even if the model exists.
+ proxies (Optional[Dict]): Proxy configuration for downloads.
+ resume_download (bool): Resume interrupted downloads.
+ local_files_only (bool): Use only local files, don't download.
+ token (Union[str, bool, None]): Token for API authentication.
+ map_location (str): Device to map model to. Defaults to "cpu".
+ strict (bool): Enforce strict state_dict loading.
+ **model_kwargs: Additional keyword arguments for model initialization.
+
+ Returns:
+ An instance of the model loaded from the pretrained weights.
+ """
+ model_dir = Path(model_id)
+ if not model_dir.exists():
+ model_dir = Path(
+ snapshot_download(
+ repo_id=model_id,
+ revision=revision,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ resume_download=resume_download,
+ token=token,
+ local_files_only=local_files_only,
+ )
+ )
+
+ config_file = model_dir / "config.json"
+ with open(config_file, 'r') as f:
+ config = json.load(f)
+
+ latent_size = config.get('latent_size')
+ if latent_size is None:
+ raise ValueError("The configuration file is missing the 'latent_size' key.")
+
+ model = cls(latent_size, **model_kwargs)
+
+ model_file = model_dir / "model_conv_vae_256_epoch_304.pth"
+ if not model_file.exists():
+ raise FileNotFoundError(f"The model checkpoint '{model_file}' does not exist.")
+ 
+ state_dict = torch.load(model_file, map_location=map_location)
+
+ new_state_dict = {}
+ for k, v in state_dict.items():
+ if k.startswith('_orig_mod.'):
+ new_state_dict[k[len('_orig_mod.'):]] = v
+ else:
+ new_state_dict[k] = v
+
+ model.load_state_dict(new_state_dict, strict=strict)
+ model.to(map_location)
+
+ return model
+
+
+model = ConvVAE.from_pretrained(
+ model_id="BioMike/classical_portrait_vae",
+ cache_dir="./model_cache",
+ map_location="cpu",
+ strict=True).eval()