Update utils.py

utils.py

@@ -5,7 +5,6 @@ import torch
 from huggan.pytorch.lightweight_gan.lightweight_gan import LightweightGAN
 from huggingface_hub import hf_hub_download
 
-model_id = "ceyda/butterfly_cropped_uniq1K_512"
 CONFIG_NAME = "config.json"
 revision = None
 cache_dir = None
@@ -16,39 +15,7 @@ local_files_only = False
 token = None
 
 
-# Load the config
-config_file = hf_hub_download(
-    repo_id=str(model_id),
-    filename=CONFIG_NAME,
-    revision=revision,
-    cache_dir=cache_dir,
-    force_download=force_download,
-    proxies=proxies,
-    resume_download=resume_download,
-    token=token,
-    local_files_only=local_files_only,
-)
-with open(config_file, "r", encoding="utf-8") as f:
-    config = json.load(f)
-
-# Call the _from_pretrained with all the needed arguments
-gan = LightweightGAN(latent_dim=256, image_size=512)
-
-gan = gan._from_pretrained(
-    model_id=str(model_id),
-    revision=revision,
-    cache_dir=cache_dir,
-    force_download=force_download,
-    proxies=proxies,
-    resume_download=resume_download,
-    local_files_only=local_files_only,
-    token=token,
-    use_auth_token=False,
-    config=config,  # usually in **model_kwargs
-)
-
-
-def carga_modelo(model_name=model_id, model_version=None):
+def carga_modelo(model_name="ceyda/butterfly_cropped_uniq1K_512"):
     """
     Loads a pre-trained LightweightGAN model from Hugging Face Model Hub.
 
@@ -59,6 +26,37 @@ def carga_modelo(model_name=model_id, model_version=None):
     Returns:
         LightweightGAN: The loaded pre-trained model.
     """
+    # Load the config
+    config_file = hf_hub_download(
+        repo_id=str(model_name),
+        filename=CONFIG_NAME,
+        revision=revision,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        proxies=proxies,
+        resume_download=resume_download,
+        token=token,
+        local_files_only=local_files_only,
+    )
+    with open(config_file, "r", encoding="utf-8") as f:
+        config = json.load(f)
+
+    # Call the _from_pretrained with all the needed arguments
+    gan = LightweightGAN(latent_dim=256, image_size=512)
+
+    gan = gan._from_pretrained(
+        model_id=str(model_name),
+        revision=revision,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        proxies=proxies,
+        resume_download=resume_download,
+        local_files_only=local_files_only,
+        token=token,
+        use_auth_token=False,
+        config=config,  # usually in **model_kwargs
+    )
+
     gan.eval()
     return gan
 
@@ -78,4 +76,3 @@ def genera(gan, batch_size=1):
         ims = gan.G(torch.randn(batch_size, gan.latent_dim)).clamp_(0.0, 1.0) * 255
         ims = ims.permute(0, 2, 3, 1).detach().cpu().numpy().astype(np.uint8)
     return ims
-