From 6f0a91f777790e613ce326726b21845d124ec8ac Mon Sep 17 00:00:00 2001 From: nachoDRT Date: Sun, 28 Sep 2025 17:15:01 +0200 Subject: [PATCH 1/5] Dockerized Version --- README.md | 16 ++++++++ docker_src/docker_start.sh | 3 ++ docker_src/docker_test.py | 55 ++++++++++++++++++++++++++++ dockerfiles/rtx4090/Dockerfile | 32 ++++++++++++++++ dockerfiles/rtx4090/requirements.txt | 4 ++ 5 files changed, 110 insertions(+) create mode 100644 docker_src/docker_start.sh create mode 100644 docker_src/docker_test.py create mode 100644 dockerfiles/rtx4090/Dockerfile create mode 100644 dockerfiles/rtx4090/requirements.txt diff --git a/README.md b/README.md index 0abb187..52841dc 100644 --- a/README.md +++ b/README.md @@ -109,3 +109,19 @@ sample( ) ``` Visualization results are stored at `output/jet.png` + +### Run docker + +Execute from ./jet-pytorch/ + +Available cards: `RTX4090` + +#### Create the docker +```bash +docker build -f dockerfiles/your_card/Dockerfile -t jet . +``` + +#### Run the docker +```bash +docker run -it --gpus '"device=0"' -v $(pwd)/output:/app/output --ipc=host jet +``` \ No newline at end of file diff --git a/docker_src/docker_start.sh b/docker_src/docker_start.sh new file mode 100644 index 0000000..f042a6e --- /dev/null +++ b/docker_src/docker_start.sh @@ -0,0 +1,3 @@ +chmod +x docker_src/docker_test.py + +python docker_src/docker_test.py \ diff --git a/docker_src/docker_test.py b/docker_src/docker_test.py new file mode 100644 index 0000000..44a5a3a --- /dev/null +++ b/docker_src/docker_test.py @@ -0,0 +1,55 @@ +from torch.distributions import Normal +from jet_pytorch import Jet +from torchvision.utils import save_image +import os +from jet_pytorch.util import get_pretrained + +# Config Jet model +jet_config = dict( + patch_size=4, + patch_dim=48, + n_patches=256, + coupling_layers=32, + block_depth=2, + block_width=512, + num_heads=8, + scale_factor=2.0, + coupling_types=( + "channels", "channels", + "channels", "channels", + "spatial", + ), + spatial_coupling_projs=( + "checkerboard", "checkerboard-inv", + "vstripes", "vstripes-inv", + "hstripes", "hstripes-inv", + ) +) + +model = Jet(**jet_config) +weights = get_pretrained() +model.load_state_dict(weights) + +# Generate latent vectors z ~ N(0, 1) +batch_size = 16 +n_patches = 256 +patch_dim = 48 +z = Normal(0, 1).sample((batch_size, n_patches, patch_dim)) + +# Rebuild imgs from z +img, _ = model.inverse(z) # img: (B, H, W, C) + +# Reformat (B, C, H, W) +img = img.permute(0, 3, 1, 2) + +# Normalize imgs [0,1] +img_min = img.amin(dim=(1, 2, 3), keepdim=True) +img_max = img.amax(dim=(1, 2, 3), keepdim=True) +img = (img - img_min) / (img_max - img_min + 1e-6) +img = img.clamp(0, 1) + +# Save imgs +os.makedirs("output", exist_ok=True) +for i in range(img.size(0)): + save_image(img[i], f"output/sample_{i:03d}.png") +save_image(img, "output/grid.png", nrow=4) diff --git a/dockerfiles/rtx4090/Dockerfile b/dockerfiles/rtx4090/Dockerfile new file mode 100644 index 0000000..75f875e --- /dev/null +++ b/dockerfiles/rtx4090/Dockerfile @@ -0,0 +1,32 @@ +FROM nvidia/cuda:11.8.0-base-ubuntu22.04 + +RUN apt-get update && apt-get install -y \ + python3.10 python3-pip libjpeg-dev zlib1g-dev git \ + && rm -rf /var/lib/apt/lists/* + +RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 +WORKDIR /app + +RUN pip install --upgrade pip + +RUN pip install --no-cache-dir "numpy<2" + +RUN pip install --no-cache-dir \ + --extra-index-url https://download.pytorch.org/whl/cu118 \ + torch==2.1.2+cu118 \ + torchvision==0.16.2+cu118 + +COPY dockerfiles/rtx4090/requirements.txt /tmp/requirements.txt +RUN pip install --no-cache-dir -r /tmp/requirements.txt + +RUN pip check + +RUN ln -sf /usr/bin/python3.10 /usr/bin/python + + +COPY . . + +ENV PYTHONPATH=/app + +RUN chmod +x docker_src/docker_start.sh +CMD ["bash", "docker_src/docker_start.sh"] diff --git a/dockerfiles/rtx4090/requirements.txt b/dockerfiles/rtx4090/requirements.txt new file mode 100644 index 0000000..9059557 --- /dev/null +++ b/dockerfiles/rtx4090/requirements.txt @@ -0,0 +1,4 @@ +einops +fire +safetensors +huggingface-hub==0.14.0 \ No newline at end of file From 1a2474ef2d52c7d4b3ce4d183126cfab54de5ce5 Mon Sep 17 00:00:00 2001 From: nachoDRT Date: Sun, 28 Sep 2025 17:42:41 +0200 Subject: [PATCH 2/5] Train in Docker with MNIST Dataset from HuggingFace --- docker_src/docker_start.sh | 5 +- docker_src/docker_train.py | 85 ++++++++++++++++++++++++++++ dockerfiles/rtx4090/requirements.txt | 3 +- 3 files changed, 91 insertions(+), 2 deletions(-) create mode 100644 docker_src/docker_train.py diff --git a/docker_src/docker_start.sh b/docker_src/docker_start.sh index f042a6e..544feef 100644 --- a/docker_src/docker_start.sh +++ b/docker_src/docker_start.sh @@ -1,3 +1,6 @@ chmod +x docker_src/docker_test.py +chmod +x docker_src/docker_train.py -python docker_src/docker_test.py \ +# python docker_src/docker_test.py \ +python docker_src/docker_train.py \ + --dataset_name "ylecun/mnist" \ No newline at end of file diff --git a/docker_src/docker_train.py b/docker_src/docker_train.py new file mode 100644 index 0000000..0086cb9 --- /dev/null +++ b/docker_src/docker_train.py @@ -0,0 +1,85 @@ +import torchvision.io +from torchvision.io import read_image as original_read_image + +def safe_read_image(path): + return original_read_image(str(path)) + +# Monkey patch global +import torchvision +torchvision.io.read_image = safe_read_image + +import argparse +import os +from tqdm import tqdm +from jet_pytorch.train import train +from datasets import load_dataset + + +def save_hf_dataset_to_disk(hf_dataset, output_dir, percentage=0.1): + os.makedirs(output_dir, exist_ok=True) + subset = hf_dataset.select(range(int(len(hf_dataset) * percentage))) + for i, sample in tqdm(enumerate(subset), total=len(subset)): + img = sample["image"] + img = sample["image"].convert("RGB").resize((64, 64)) + img.save(os.path.join(output_dir, f"{i}.png")) + + +def get_hf_dataset(): + + dataset = load_dataset(dataset_name) + save_hf_dataset_to_disk(dataset["train"], f"./{dataset_name}_train") + save_hf_dataset_to_disk(dataset["test"], f"./{dataset_name}_valid") + + +def train_jet(): + + jet_config = dict( + patch_size=4, + patch_dim=48, + n_patches=256, + coupling_layers=32, + block_depth=2, + block_width=512, + num_heads=8, + scale_factor=2.0, + coupling_types=( + "channels", "channels", + "channels", "channels", + "spatial", + ), + spatial_coupling_projs=( + "checkerboard", "checkerboard-inv", + "vstripes", "vstripes-inv", + "hstripes", "hstripes-inv", + ) + ) + + + train( + jet_config=jet_config, + batch_size=32, + accumulate_steps=16, + device="cuda:0", + epochs=50, + warmup_percentage=0.1, + max_grad_norm=1.0, + learning_rate=3e-4, + weight_decay=1e-5, + adam_betas=(0.9, 0.95), + images_path_train=f"./{dataset_name}_train", + images_path_valid= f"./{dataset_name}_valid", + num_workers=8, + checkpoint_path="jet.pt", + ) + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument("--dataset_name", type=str) + args = parser.parse_args() + + dataset_name = args.dataset_name + + get_hf_dataset() + train_jet() diff --git a/dockerfiles/rtx4090/requirements.txt b/dockerfiles/rtx4090/requirements.txt index 9059557..d935784 100644 --- a/dockerfiles/rtx4090/requirements.txt +++ b/dockerfiles/rtx4090/requirements.txt @@ -1,4 +1,5 @@ einops fire safetensors -huggingface-hub==0.14.0 \ No newline at end of file +huggingface-hub==0.14.0 +datasets \ No newline at end of file From 010b0dbae0a56fd7ccb29fab5749c7c1be57bdaa Mon Sep 17 00:00:00 2001 From: nachoDRT Date: Sun, 28 Sep 2025 18:26:47 +0200 Subject: [PATCH 3/5] Track Training Metrics in WandB --- docker_src/docker_start.sh | 5 ++++- docker_src/docker_train.py | 27 ++++++++++++++++++++++++++- dockerfiles/rtx4090/Dockerfile | 1 + dockerfiles/rtx4090/requirements.txt | 3 ++- jet_pytorch/train.py | 15 +++++++++++++++ 5 files changed, 48 insertions(+), 3 deletions(-) diff --git a/docker_src/docker_start.sh b/docker_src/docker_start.sh index 544feef..673bd08 100644 --- a/docker_src/docker_start.sh +++ b/docker_src/docker_start.sh @@ -3,4 +3,7 @@ chmod +x docker_src/docker_train.py # python docker_src/docker_test.py \ python docker_src/docker_train.py \ - --dataset_name "ylecun/mnist" \ No newline at end of file + --dataset_name "ylecun/mnist" \ + --wandb_entity "ciclab-comillas" \ + --wandb_project "jet" \ + --wandb_run_name "mnist" \ No newline at end of file diff --git a/docker_src/docker_train.py b/docker_src/docker_train.py index 0086cb9..bd58513 100644 --- a/docker_src/docker_train.py +++ b/docker_src/docker_train.py @@ -1,6 +1,7 @@ import torchvision.io from torchvision.io import read_image as original_read_image +# TODO Fix this properly def safe_read_image(path): return original_read_image(str(path)) @@ -13,9 +14,10 @@ def safe_read_image(path): from tqdm import tqdm from jet_pytorch.train import train from datasets import load_dataset +import wandb -def save_hf_dataset_to_disk(hf_dataset, output_dir, percentage=0.1): +def save_hf_dataset_to_disk(hf_dataset, output_dir, percentage=0.01): os.makedirs(output_dir, exist_ok=True) subset = hf_dataset.select(range(int(len(hf_dataset) * percentage))) for i, sample in tqdm(enumerate(subset), total=len(subset)): @@ -33,6 +35,23 @@ def get_hf_dataset(): def train_jet(): + wandb.login(key=os.getenv("WANDB_API_KEY")) + wandb.init( + project=wandb_project, + name=wandb_run_name, + entity=wandb_entity, + config={ + "batch_size": 64, + "accumulate_steps": 16, + "epochs": 50, + "learning_rate": 3e-4, + "patch_size": 4, + "patch_dim": 48, + "n_patches": 256, + "coupling_layers": 32, + }, + ) + jet_config = dict( patch_size=4, patch_dim=48, @@ -77,9 +96,15 @@ def train_jet(): parser = argparse.ArgumentParser() parser.add_argument("--dataset_name", type=str) + parser.add_argument("--wandb_entity", type=str) + parser.add_argument("--wandb_project", type=str) + parser.add_argument("--wandb_run_name", type=str) args = parser.parse_args() dataset_name = args.dataset_name + wandb_entity = args.wandb_entity + wandb_project = args.wandb_project + wandb_run_name = args.wandb_run_name get_hf_dataset() train_jet() diff --git a/dockerfiles/rtx4090/Dockerfile b/dockerfiles/rtx4090/Dockerfile index 75f875e..a71d462 100644 --- a/dockerfiles/rtx4090/Dockerfile +++ b/dockerfiles/rtx4090/Dockerfile @@ -23,6 +23,7 @@ RUN pip check RUN ln -sf /usr/bin/python3.10 /usr/bin/python +ENV WANDB_API_KEY= COPY . . diff --git a/dockerfiles/rtx4090/requirements.txt b/dockerfiles/rtx4090/requirements.txt index d935784..5f80325 100644 --- a/dockerfiles/rtx4090/requirements.txt +++ b/dockerfiles/rtx4090/requirements.txt @@ -2,4 +2,5 @@ einops fire safetensors huggingface-hub==0.14.0 -datasets \ No newline at end of file +datasets +wandb \ No newline at end of file diff --git a/jet_pytorch/train.py b/jet_pytorch/train.py index 40cc96f..9d1c116 100644 --- a/jet_pytorch/train.py +++ b/jet_pytorch/train.py @@ -7,6 +7,7 @@ from torch.utils.data import DataLoader from torchvision.transforms import v2 from tqdm import tqdm +import wandb from jet_pytorch import Jet from jet_pytorch.util import bits_per_dim @@ -121,6 +122,14 @@ def train( gn=gn, lr=lr, )) + wandb.log({ + "train/bpd": log_bpd, + "train/nll": log_nll, + "train/logdet": log_logdet, + "train/grad_norm": gn, + "train/lr": lr, + "epoch": epoch, + }) steps = 0 log_bpd = 0.0 log_nll = 0.0 @@ -159,6 +168,12 @@ def train( val_nll /= len(val_dataloader) val_logdet /= len(val_dataloader) print(f"validation metrics - bpd: {val_bpd:.2f}, nll: {val_nll:.2f}, logdet: {val_logdet:.2f}") + wandb.log({ + "val/bpd": val_bpd, + "val/nll": val_nll, + "val/logdet": val_logdet, + "epoch": epoch, + }) model.train() From 78408c2ee55e8f8be743d7386983b16d67403b37 Mon Sep 17 00:00:00 2001 From: nachoDRT Date: Sun, 28 Sep 2025 23:51:23 +0200 Subject: [PATCH 4/5] Iterativetily Save Best Model after Validation in HuggingFace Repo --- docker_src/docker_start.sh | 3 ++- docker_src/docker_train.py | 8 +++++++- dockerfiles/rtx4090/Dockerfile | 1 + jet_pytorch/train.py | 21 +++++++++++++++++++++ 4 files changed, 31 insertions(+), 2 deletions(-) diff --git a/docker_src/docker_start.sh b/docker_src/docker_start.sh index 673bd08..21c14b4 100644 --- a/docker_src/docker_start.sh +++ b/docker_src/docker_start.sh @@ -6,4 +6,5 @@ python docker_src/docker_train.py \ --dataset_name "ylecun/mnist" \ --wandb_entity "ciclab-comillas" \ --wandb_project "jet" \ - --wandb_run_name "mnist" \ No newline at end of file + --wandb_run_name "mnist" \ + --hf_repo_id "your_hf_repo" \ No newline at end of file diff --git a/docker_src/docker_train.py b/docker_src/docker_train.py index bd58513..a58dc1b 100644 --- a/docker_src/docker_train.py +++ b/docker_src/docker_train.py @@ -15,6 +15,7 @@ def safe_read_image(path): from jet_pytorch.train import train from datasets import load_dataset import wandb +from huggingface_hub import login def save_hf_dataset_to_disk(hf_dataset, output_dir, percentage=0.01): @@ -52,6 +53,8 @@ def train_jet(): }, ) + login(token=os.getenv("HUGGINGFACE_HUB_TOKEN")) + jet_config = dict( patch_size=4, patch_dim=48, @@ -79,7 +82,7 @@ def train_jet(): batch_size=32, accumulate_steps=16, device="cuda:0", - epochs=50, + epochs=250, warmup_percentage=0.1, max_grad_norm=1.0, learning_rate=3e-4, @@ -89,6 +92,7 @@ def train_jet(): images_path_valid= f"./{dataset_name}_valid", num_workers=8, checkpoint_path="jet.pt", + hf_repo_id=hf_repo_id ) @@ -99,12 +103,14 @@ def train_jet(): parser.add_argument("--wandb_entity", type=str) parser.add_argument("--wandb_project", type=str) parser.add_argument("--wandb_run_name", type=str) + parser.add_argument("--hf_repo_id", type=str, default=None) args = parser.parse_args() dataset_name = args.dataset_name wandb_entity = args.wandb_entity wandb_project = args.wandb_project wandb_run_name = args.wandb_run_name + hf_repo_id = args.hf_repo_id get_hf_dataset() train_jet() diff --git a/dockerfiles/rtx4090/Dockerfile b/dockerfiles/rtx4090/Dockerfile index a71d462..6d8d24e 100644 --- a/dockerfiles/rtx4090/Dockerfile +++ b/dockerfiles/rtx4090/Dockerfile @@ -24,6 +24,7 @@ RUN pip check RUN ln -sf /usr/bin/python3.10 /usr/bin/python ENV WANDB_API_KEY= +ENV HUGGINGFACE_HUB_TOKEN= COPY . . diff --git a/jet_pytorch/train.py b/jet_pytorch/train.py index 9d1c116..f9f0168 100644 --- a/jet_pytorch/train.py +++ b/jet_pytorch/train.py @@ -8,6 +8,8 @@ from torchvision.transforms import v2 from tqdm import tqdm import wandb +from huggingface_hub import upload_file +from safetensors.torch import save_file from jet_pytorch import Jet from jet_pytorch.util import bits_per_dim @@ -31,6 +33,7 @@ def train( num_workers=16, device="cuda:0", checkpoint_path="jet_imagenet.pt", + hf_repo_id = None ): t.set_float32_matmul_precision("medium") @@ -88,6 +91,8 @@ def train( log_nll = 0.0 log_logdet = 0.0 + best_val_loss = float("inf") + for epoch in range(epochs): pbar = tqdm(dataloader, total=len(dataloader) // accumulate_steps) @@ -174,6 +179,22 @@ def train( "val/logdet": val_logdet, "epoch": epoch, }) + + # Save model if it's the best so far + if val_bpd < best_val_loss: + print(f"New best validation loss: {val_bpd:.4f} (prev {best_val_loss:.4f})") + best_val_loss = val_bpd + + state_dict = orig_model.state_dict() + save_file(state_dict, "jet_mnist.safetensors") + + upload_file( + path_or_fileobj="jet_mnist.safetensors", + path_in_repo="jet_mnist.safetensors", + repo_id=hf_repo_id, + repo_type="model" + ) + model.train() From 03de5a1e02217d3ca78b190140b30ed352a04468 Mon Sep 17 00:00:00 2001 From: nachoDRT Date: Tue, 30 Sep 2025 18:48:13 +0200 Subject: [PATCH 5/5] Test with Custom Trained Model --- .gitignore | 1 + docker_src/docker_start.sh | 17 +++--- docker_src/docker_test.py | 111 +++++++++++++++++++++---------------- jet_pytorch/util.py | 14 +++-- 4 files changed, 83 insertions(+), 60 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..77320b3 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +output/* diff --git a/docker_src/docker_start.sh b/docker_src/docker_start.sh index 21c14b4..c80fa32 100644 --- a/docker_src/docker_start.sh +++ b/docker_src/docker_start.sh @@ -1,10 +1,13 @@ chmod +x docker_src/docker_test.py chmod +x docker_src/docker_train.py -# python docker_src/docker_test.py \ -python docker_src/docker_train.py \ - --dataset_name "ylecun/mnist" \ - --wandb_entity "ciclab-comillas" \ - --wandb_project "jet" \ - --wandb_run_name "mnist" \ - --hf_repo_id "your_hf_repo" \ No newline at end of file +python docker_src/docker_test.py \ + --repo_id "de-Rodrigo/jet-mnist" \ + --model_name "jet_mnist" + +# python docker_src/docker_train.py \ +# --dataset_name "ylecun/mnist" \ +# --wandb_entity "ciclab-comillas" \ +# --wandb_project "jet" \ +# --wandb_run_name "mnist" \ +# --hf_repo_id "de-Rodrigo/jet-mnist" \ No newline at end of file diff --git a/docker_src/docker_test.py b/docker_src/docker_test.py index 44a5a3a..d144303 100644 --- a/docker_src/docker_test.py +++ b/docker_src/docker_test.py @@ -3,53 +3,68 @@ from torchvision.utils import save_image import os from jet_pytorch.util import get_pretrained +import argparse +from huggingface_hub import login -# Config Jet model -jet_config = dict( - patch_size=4, - patch_dim=48, - n_patches=256, - coupling_layers=32, - block_depth=2, - block_width=512, - num_heads=8, - scale_factor=2.0, - coupling_types=( - "channels", "channels", - "channels", "channels", - "spatial", - ), - spatial_coupling_projs=( - "checkerboard", "checkerboard-inv", - "vstripes", "vstripes-inv", - "hstripes", "hstripes-inv", + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument("--repo_id", type=str, default=None) + parser.add_argument("--model_name", type=str, default=None) + args = parser.parse_args() + + repo_id = args.repo_id + model_name = args.model_name + + login(token=os.getenv("HUGGINGFACE_HUB_TOKEN")) + + # Config Jet model + jet_config = dict( + patch_size=4, + patch_dim=48, + n_patches=256, + coupling_layers=32, + block_depth=2, + block_width=512, + num_heads=8, + scale_factor=2.0, + coupling_types=( + "channels", "channels", + "channels", "channels", + "spatial", + ), + spatial_coupling_projs=( + "checkerboard", "checkerboard-inv", + "vstripes", "vstripes-inv", + "hstripes", "hstripes-inv", + ) ) -) - -model = Jet(**jet_config) -weights = get_pretrained() -model.load_state_dict(weights) - -# Generate latent vectors z ~ N(0, 1) -batch_size = 16 -n_patches = 256 -patch_dim = 48 -z = Normal(0, 1).sample((batch_size, n_patches, patch_dim)) - -# Rebuild imgs from z -img, _ = model.inverse(z) # img: (B, H, W, C) - -# Reformat (B, C, H, W) -img = img.permute(0, 3, 1, 2) - -# Normalize imgs [0,1] -img_min = img.amin(dim=(1, 2, 3), keepdim=True) -img_max = img.amax(dim=(1, 2, 3), keepdim=True) -img = (img - img_min) / (img_max - img_min + 1e-6) -img = img.clamp(0, 1) - -# Save imgs -os.makedirs("output", exist_ok=True) -for i in range(img.size(0)): - save_image(img[i], f"output/sample_{i:03d}.png") -save_image(img, "output/grid.png", nrow=4) + + model = Jet(**jet_config) + weights = get_pretrained(repo_id, model_name) + model.load_state_dict(weights) + + # Generate latent vectors z ~ N(0, 1) + batch_size = 16 + n_patches = 256 + patch_dim = 48 + z = Normal(0, 1).sample((batch_size, n_patches, patch_dim)) + + # Rebuild imgs from z + img, _ = model.inverse(z) # img: (B, H, W, C) + + # Reformat (B, C, H, W) + img = img.permute(0, 3, 1, 2) + + # Normalize imgs [0,1] + img_min = img.amin(dim=(1, 2, 3), keepdim=True) + img_max = img.amax(dim=(1, 2, 3), keepdim=True) + img = (img - img_min) / (img_max - img_min + 1e-6) + img = img.clamp(0, 1) + + # Save imgs + os.makedirs("output", exist_ok=True) + for i in range(img.size(0)): + save_image(img[i], f"output/sample_{i:03d}.png") + save_image(img, "output/grid.png", nrow=4) diff --git a/jet_pytorch/util.py b/jet_pytorch/util.py index 90d0333..b0d6c37 100644 --- a/jet_pytorch/util.py +++ b/jet_pytorch/util.py @@ -12,16 +12,20 @@ from torchvision.io import read_image -def get_pretrained(): - pretrained_path = "models/jet_imagenet64x64_200m/jet_imagenet64x64_200m.safetensors" +def get_pretrained(repo_id: str = None, model_name:str = None): + if model_name is None: + pretrained_path = "models/jet_imagenet64x64_200m/jet_imagenet64x64_200m.safetensors" + else: + pretrained_path = f"models/{model_name}/{model_name}.safetensors" os.makedirs("models", exist_ok=True) if not os.path.exists(pretrained_path): _ = hf_hub_download( - repo_id="btrude/jet_imagenet64x64_200m", - filename="jet_imagenet64x64_200m.safetensors", - local_dir="models/jet_imagenet64x64_200m", + repo_id=repo_id, + filename=f"{model_name}.safetensors", + local_dir=f"models/{model_name}", ) + return load_file(pretrained_path)