Spaces:
Runtime error
Runtime error
import argparse | |
import os | |
import shutil | |
from transformers import ( | |
BlipForConditionalGeneration, | |
BlipProcessor, | |
CLIPSegForImageSegmentation, | |
CLIPSegProcessor, | |
Swin2SRForImageSuperResolution, | |
) | |
DEFAULT_BLIP = "Salesforce/blip-image-captioning-large" | |
DEFAULT_CLIPSEG = "CIDAS/clipseg-rd64-refined" | |
DEFAULT_SWINIR = "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr" | |
def upload(args): | |
blip_processor = BlipProcessor.from_pretrained(DEFAULT_BLIP) | |
blip_model = BlipForConditionalGeneration.from_pretrained(DEFAULT_BLIP) | |
clip_processor = CLIPSegProcessor.from_pretrained(DEFAULT_CLIPSEG) | |
clip_model = CLIPSegForImageSegmentation.from_pretrained(DEFAULT_CLIPSEG) | |
swin_model = Swin2SRForImageSuperResolution.from_pretrained(DEFAULT_SWINIR) | |
temp_models = "tmp/models" | |
if os.path.exists(temp_models): | |
shutil.rmtree(temp_models) | |
os.makedirs(temp_models) | |
blip_processor.save_pretrained(os.path.join(temp_models, "blip_processor")) | |
blip_model.save_pretrained(os.path.join(temp_models, "blip_large")) | |
clip_processor.save_pretrained(os.path.join(temp_models, "clip_seg_processor")) | |
clip_model.save_pretrained(os.path.join(temp_models, "clip_seg_rd64_refined")) | |
swin_model.save_pretrained( | |
os.path.join(temp_models, "swin2sr_realworld_sr_x4_64_bsrgan_psnr") | |
) | |
for val in os.listdir(temp_models): | |
if "tar" not in val: | |
os.system( | |
f"sudo tar -cvf {os.path.join(temp_models, val)}.tar -C {os.path.join(temp_models, val)} ." | |
) | |
os.system( | |
f"gcloud storage cp -R {os.path.join(temp_models, val)}.tar gs://{args.bucket}/{val}/" | |
) | |
if __name__ == "__main__": | |
parser = argparse.ArgumentParser() | |
parser.add_argument("--bucket", "-m", type=str) | |
args = parser.parse_args() | |
upload(args) | |