-
Notifications
You must be signed in to change notification settings - Fork 4
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
🤗 Add support for downloading Huggingface weights. #2
Changes from 2 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,54 @@ | ||
import cv2 | ||
import math | ||
import sys | ||
import torch | ||
import numpy as np | ||
import argparse | ||
from imageio import mimsave | ||
|
||
'''==========import from our code==========''' | ||
sys.path.append('.') | ||
import config as cfg | ||
from Trainer_finetune import Model | ||
from benchmark.utils.padder import InputPadder | ||
|
||
parser = argparse.ArgumentParser() | ||
parser.add_argument('--model', default='VFIMamba_S', type=str) | ||
parser.add_argument('--scale', default=0, type=float) | ||
|
||
args = parser.parse_args() | ||
assert args.model in ['VFIMamba_S', 'VFIMamba'], 'Model not exists!' | ||
|
||
|
||
'''==========Model setting==========''' | ||
TTA = False | ||
if args.model == 'VFIMamba': | ||
TTA = True | ||
cfg.MODEL_CONFIG['LOGNAME'] = 'VFIMamba' | ||
cfg.MODEL_CONFIG['MODEL_ARCH'] = cfg.init_model_config( | ||
F = 32, | ||
depth = [2, 2, 2, 3, 3] | ||
) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Could these config values be stored in a (just a suggestion, I'm not entirely sure how There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Absolutely! Thanks to your advice, we can now directly initialize a model when loading weights from Hugging Face using the following code (as also demonstrated in hf_demo_2x.py): model = Model.from_pretrained(args.model) Cheers! There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Awesome! |
||
model = Model(-1) | ||
model.from_pretrained(args.model) | ||
model.eval() | ||
model.device() | ||
|
||
|
||
print(f'=========================Start Generating=========================') | ||
|
||
I0 = cv2.imread('example/im1.png') | ||
I2 = cv2.imread('example/im2.png') | ||
|
||
I0_ = (torch.tensor(I0.transpose(2, 0, 1)).cuda() / 255.).unsqueeze(0) | ||
I2_ = (torch.tensor(I2.transpose(2, 0, 1)).cuda() / 255.).unsqueeze(0) | ||
|
||
padder = InputPadder(I0_.shape, divisor=32) | ||
I0_, I2_ = padder.pad(I0_, I2_) | ||
|
||
mid = (padder.unpad(model.inference(I0_, I2_, True, TTA=TTA, fast_TTA=TTA, scale=args.scale))[0].detach().cpu().numpy().transpose(1, 2, 0) * 255.0).astype(np.uint8) | ||
images = [I0[:, :, ::-1], mid[:, :, ::-1], I2[:, :, ::-1]] | ||
mimsave('example/out_2x_hf.gif', images, fps=3) | ||
|
||
|
||
print(f'=========================Done=========================') |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Hi there, maintainer of
huggingface_hub
here 👋 May I suggest a more opinionated implementation for this method? In particular:huggingface_hub
required and raise an explicit error if it's not the case. This way you don't have to maintain different ways to load the model. In particular, end users will benefit from thehuggingface_hub
cache when running the script several times.model.pkl
in both cases. This way, the repo id and filename don't have to be correlated. Also, it would simplify the way downloads are counted.model_id
as input, not only the 2 official one from the repo. This way, your library can be reused by anyone wanting the train, retrain, fine-tuned, etc new models based on your work. This should greatly improve adoption of the library.model_id
can be either a full repo id (e.g.MCG-NJU/VFIMamba_S
) or simply a model name in which case theMCG-NJU/
org is prefixed.from_pretrained
a class method to initialize the object directly withmodel = Model.from_pretrained("VFIMamba_S", local_rank=-1)
Those changes are not an obligation to make things work with the Hub but in my opinion it would greatly improve the integration. Let me know what you think!
(cc @hanouticelina for viz')
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you very much! Your suggestion is very helpful. I've made some changes based on your advice.