Update clip

app.py

@@ -15,33 +15,6 @@ from libs.muse import MUSE
 import utils
 import numpy as np
 from PIL import Image
-print("cuda available:",torch.cuda.is_available())
-print("cuda device count:",torch.cuda.device_count())
-print("cuda device name:",torch.cuda.get_device_name(0))
-# print(os.system("nvidia-smi"))
-print(os.system("nvcc --version"))
-
-empty_context = np.load("assets/contexts/empty_context.npy")
-
-print("downloading cc3m-285000.ckpt")
-os.makedirs("assets/ckpts/cc3m-285000.ckpt",exist_ok=True)
-
-wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/lr_scheduler.pth","assets/ckpts/cc3m-285000.ckpt/lr_scheduler.pth")
-wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/optimizer.pth","assets/ckpts/cc3m-285000.ckpt/optimizer.pth")
-wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/nnet.pth","assets/ckpts/cc3m-285000.ckpt/nnet.pth")
-wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/nnet_ema.pth","assets/ckpts/cc3m-285000.ckpt/nnet_ema.pth")
-wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/step.pth","assets/ckpts/cc3m-285000.ckpt/step.pth")
-wget.download("https://huggingface.co/zideliu/vqgan/resolve/main/vqgan_jax_strongaug.ckpt","assets/vqgan_jax_strongaug.ckpt")
-os.system("ls assets/ckpts/cc3m-285000.ckpt")
-def set_seed(seed: int):
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    
-def d(**kwargs):
-    """Helper of creating a config dict."""
-    return ml_collections.ConfigDict(initial_dictionary=kwargs)
 
 def get_config():
     config = ml_collections.ConfigDict()
@@ -98,6 +71,50 @@ def get_config():
     )
     return config
 
+print("cuda available:",torch.cuda.is_available())
+print("cuda device count:",torch.cuda.device_count())
+print("cuda device name:",torch.cuda.get_device_name(0))
+# print(os.system("nvidia-smi"))
+print(os.system("nvcc --version"))
+
+empty_context = np.load("assets/contexts/empty_context.npy")
+
+config = get_config()
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+print(device)
+# Load open_clip and vq model
+print("GPU memory:",torch.cuda.memory_allocated(0)/1024/1024/1024,"GB")
+prompt_model,_,_ = open_clip.create_model_and_transforms('ViT-bigG-14', 'laion2b_s39b_b160k',device=device)
+print("GPU memory:",torch.cuda.memory_allocated(0)/1024/1024/1024,"GB")
+
+prompt_model = prompt_model.to(device)
+prompt_model.eval()
+tokenizer = open_clip.get_tokenizer('ViT-bigG-14')
+
+
+
+
+print("downloading cc3m-285000.ckpt")
+os.makedirs("assets/ckpts/cc3m-285000.ckpt",exist_ok=True)
+
+wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/lr_scheduler.pth","assets/ckpts/cc3m-285000.ckpt/lr_scheduler.pth")
+wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/optimizer.pth","assets/ckpts/cc3m-285000.ckpt/optimizer.pth")
+wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/nnet.pth","assets/ckpts/cc3m-285000.ckpt/nnet.pth")
+wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/nnet_ema.pth","assets/ckpts/cc3m-285000.ckpt/nnet_ema.pth")
+wget.download("https://huggingface.co/nzl-thu/MUSE/resolve/main/assets/ckpts/cc3m-285000.ckpt/step.pth","assets/ckpts/cc3m-285000.ckpt/step.pth")
+wget.download("https://huggingface.co/zideliu/vqgan/resolve/main/vqgan_jax_strongaug.ckpt","assets/vqgan_jax_strongaug.ckpt")
+os.system("ls assets/ckpts/cc3m-285000.ckpt")
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    
+def d(**kwargs):
+    """Helper of creating a config dict."""
+    return ml_collections.ConfigDict(initial_dictionary=kwargs)
+
+
 def cfg_nnet(x, context, scale=None,lambdaA=None,lambdaB=None):
     _cond = nnet_ema(x, context=context)
     _cond_w_adapter = nnet_ema(x,context=context,use_adapter=True)
@@ -113,14 +130,7 @@ def unprocess(x):
     x.clamp_(0., 1.)
     return x
 
-config = get_config()
-device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-print(device)
-# Load open_clip and vq model
-prompt_model,_,_ = open_clip.create_model_and_transforms('ViT-bigG-14', 'laion2b_s39b_b160k',device='cuda')
-prompt_model = prompt_model.to(device)
-prompt_model.eval()
-tokenizer = open_clip.get_tokenizer('ViT-bigG-14')
+
 
 vq_model = taming.models.vqgan.get_model('vq-f16-jax.yaml')
 vq_model.eval()