added files

.gitignore

@@ -0,0 +1 @@
+pre_trained_checkpoint_4ch

generate_4ch.py

@@ -5,6 +5,7 @@ import torch.nn.functional as F
 from torchvision.datasets import ImageFolder
 from torch.utils.data import DataLoader
 from torchvision import utils as vutils
+from huggingface_hub import PyTorchModelHubMixin
 
 import os
 import random
@@ -36,12 +37,22 @@ def batch_save(images, folder_name):
     for i, image in enumerate(images):
         vutils.save_image(image.add(1).mul(0.5), folder_name+'/%d.jpg'%i)
 
+# To push the model to Huggingface model hub
+class MyFastGanModel(nn.Module, PyTorchModelHubMixin):
+
+    def __init__(self, config: dict) -> None:
+        super().__init__()
+
+        self.model = Generator( ngf=config["ngf"], nz=config["noise_dim"], nc=config["nc"], im_size=config["im_size"])
+
+    def forward(self, x):
+        return self.model(x)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description='generate images'
     )
-    parser.add_argument('--ckpt', type=str, default="pre_trained_checkpoint_4ch/all_50000.pth")
+    parser.add_argument('--ckpt', type=str, default="/work/vajira/DL/FastGAN-pytorch/train_results/test1_4ch/models/all_50000.pth")
     parser.add_argument('--artifacts', type=str, default=".", help='path to artifacts.')
     parser.add_argument('--cuda', type=int, default=0, help='index of gpu to use')
     parser.add_argument('--start_iter', type=int, default=6)
@@ -50,7 +61,7 @@ if __name__ == "__main__":
     parser.add_argument('--dist', type=str, default='test_out')
     parser.add_argument('--size', type=int, default=256)
     parser.add_argument('--batch', default=1, type=int, help='batch size')
-    parser.add_argument('--n_sample', type=int, default=1000)
+    parser.add_argument('--n_sample', type=int, default=1)
     parser.add_argument('--big', action='store_true')
     parser.add_argument('--im_size', type=int, default=256)
     parser.add_argument("--save_option", default="image_and_mask", help="Options to svae output, image_only, mask_only, image_and_mask", choices=["image_only","mask_only", "image_and_mask"])
@@ -59,8 +70,17 @@ if __name__ == "__main__":
 
     noise_dim = 256
     device = torch.device('cuda:%d'%(args.cuda))
+
+    # adding the model to the model hub
+    config={"ngf":64, "noise_dim":noise_dim, "nc":4, "im_size":args.im_size}
+    net_ig = MyFastGanModel(config=config)     
+
     
-    net_ig = Generator( ngf=64, nz=noise_dim, nc=4, im_size=args.im_size)#, big=args.big )
+
+    # exit
+    #exit()    
+    
+    #net_ig = model #Generator( ngf=64, nz=noise_dim, nc=4, im_size=args.im_size)#, big=args.big )
     net_ig.to(device)
 
     #for epoch in [10000*i for i in range(args.start_iter, args.end_iter+1)]:
@@ -69,13 +89,25 @@ if __name__ == "__main__":
     checkpoint = torch.load(ckpt)
     # Remove prefix `module`.
     checkpoint['g'] = {k.replace('module.', ''): v for k, v in checkpoint['g'].items()}
-    net_ig.load_state_dict(checkpoint['g'])
+    net_ig.model.load_state_dict(checkpoint['g'])
     #load_params(net_ig, checkpoint['g_ema'])
 
     #net_ig.eval()
     print("load checkpoint success")
 
     net_ig.to(device)
+    # Save locally
+    net_ig.save_pretrained("pre_trained_checkpoint_4ch", config=config)  # Save the model locally
+    print("Model saved locally. Pushing to Huggingface model hub...")
+
+    # Push to the Huggingface model hub
+    # push to the hub
+    net_ig.push_to_hub("deepsynthbody/deepfake_gi_fastGAN", config=config)
+
+
+    print("pushed to the Huggingface model hub. Done.")
+    exit()  
+
 
     del checkpoint
 

generate_4ch_from_huggingface.py

@@ -0,0 +1,160 @@
+import torch
+from torch import nn
+from torch import optim
+import torch.nn.functional as F
+from torchvision.datasets import ImageFolder
+from torch.utils.data import DataLoader
+from torchvision import utils as vutils
+from huggingface_hub import PyTorchModelHubMixin
+
+import os
+import random
+import argparse
+from tqdm import tqdm
+
+from models import Generator
+
+
+def load_params(model, new_param):
+    for p, new_p in zip(model.parameters(), new_param):
+        p.data.copy_(new_p)
+
+def resize(img):
+    return F.interpolate(img, size=256)
+
+def batch_generate(zs, netG, batch=8):
+    g_images = []
+    with torch.no_grad():
+        for i in range(len(zs)//batch):
+            g_images.append( netG(zs[i*batch:(i+1)*batch]).cpu() )
+        if len(zs)%batch>0:
+            g_images.append( netG(zs[-(len(zs)%batch):]).cpu() )
+    return torch.cat(g_images)
+
+def batch_save(images, folder_name):
+    if not os.path.exists(folder_name):
+        os.mkdir(folder_name)
+    for i, image in enumerate(images):
+        vutils.save_image(image.add(1).mul(0.5), folder_name+'/%d.jpg'%i)
+
+# To push the model to Huggingface model hub
+class MyFastGanModel(nn.Module, PyTorchModelHubMixin):
+
+    def __init__(self, config: dict) -> None:
+        super().__init__()
+
+        self.model = Generator( ngf=config["ngf"], nz=config["noise_dim"], nc=config["nc"], im_size=config["im_size"])
+
+    def forward(self, x):
+        return self.model(x)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='generate images'
+    )
+    parser.add_argument('--ckpt', type=str, default="/work/vajira/DL/FastGAN-pytorch/train_results/test1_4ch/models/all_50000.pth")
+    parser.add_argument('--artifacts', type=str, default=".", help='path to artifacts.')
+    parser.add_argument('--cuda', type=int, default=0, help='index of gpu to use')
+    parser.add_argument('--start_iter', type=int, default=6)
+    parser.add_argument('--end_iter', type=int, default=10)
+
+    parser.add_argument('--dist', type=str, default='test_out')
+    parser.add_argument('--size', type=int, default=256)
+    parser.add_argument('--batch', default=1, type=int, help='batch size')
+    parser.add_argument('--n_sample', type=int, default=1)
+    parser.add_argument('--big', action='store_true')
+    parser.add_argument('--im_size', type=int, default=256)
+    parser.add_argument("--save_option", default="image_and_mask", help="Options to svae output, image_only, mask_only, image_and_mask", choices=["image_only","mask_only", "image_and_mask"])
+    parser.set_defaults(big=False)
+    args = parser.parse_args()
+
+    noise_dim = 256
+    device = torch.device('cuda:%d'%(args.cuda))
+
+    # adding the model to the model hub
+    config={"ngf":64, "noise_dim":noise_dim, "nc":4, "im_size":args.im_size}
+    net_ig = MyFastGanModel(config=config)     
+
+    
+
+    # exit
+    #exit()    
+    
+    #net_ig = model #Generator( ngf=64, nz=noise_dim, nc=4, im_size=args.im_size)#, big=args.big )
+    #net_ig.to(device)
+
+    #for epoch in [10000*i for i in range(args.start_iter, args.end_iter+1)]:
+    #ckpt = args.ckpt #f"{args.artifacts}/models/{epoch}.pth"
+    #checkpoint = torch.load(ckpt, map_location=lambda a,b: a)
+    #checkpoint = torch.load(ckpt)
+    # Remove prefix `module`.
+    #checkpoint['g'] = {k.replace('module.', ''): v for k, v in checkpoint['g'].items()}
+    #net_ig.model.load_state_dict(checkpoint['g'])
+    #load_params(net_ig, checkpoint['g_ema'])
+
+    net_ig = MyFastGanModel.from_pretrained("deepsynthbody/deepfake_gi_fastGAN", config=config) # Load the model from the hub
+
+    #net_ig.eval()
+    print("load checkpoint success")
+
+    net_ig.to(device)
+    # Save locally
+   # net_ig.save_pretrained("pre_trained_checkpoint_4ch", config=config)  # Save the model locally
+   # print("Model saved locally. Pushing to Huggingface model hub...")
+
+    # Push to the Huggingface model hub
+    # push to the hub
+   # net_ig.push_to_hub("deepsynthbody/deepfake_gi_fastGAN", config=config)
+
+
+    #print("pushed to the Huggingface model hub. Done.")
+    #exit()  
+
+
+    #del checkpoint
+
+    #dist = 'eval_%d'%(epoch)
+    #dist = os.path.join(args.dist, 'img')
+    dist = args.dist
+    os.makedirs(dist, exist_ok=True)
+
+    with torch.no_grad():
+        for i in tqdm(range(args.n_sample//args.batch)):
+            noise = torch.randn(args.batch, noise_dim).to(device)
+            g_imgs = net_ig(noise)[0]
+            g_imgs = F.interpolate(g_imgs, 512)
+            
+            
+            for j, g_img in enumerate( g_imgs ):
+                #print("img sahpe=", g_img.shape)
+                g_mask = g_img.add(1).mul(0.5)[-1, :, :].expand(3, -1, -1)
+                g_img = g_img.add(1).mul(0.5)[0:3, :, :]
+
+                # Clean generated data using clamping
+                g_mask = torch.clamp(g_mask, min=0, max=1)
+                g_img = torch.clamp(g_img, min=0, max=1)
+                #print(g_mask.type())
+                g_mask = (g_mask > 0.5) * 1.0
+                #print(g_mask.type())
+
+                #print("gmask_min:", g_mask.min())
+                #print("gmask_max:", g_mask.max())
+                #exit()
+                
+                #print("img sahpe=", g_img.shape)
+
+                if args.save_option == "image_and_mask":
+                    vutils.save_image(g_img, 
+                        os.path.join(dist, '%d_img.png'%(i*args.batch+j)))#, normalize=True, range=(-1,1))
+                    vutils.save_image(g_mask, 
+                        os.path.join(dist, '%d_mask.png'%(i*args.batch+j))) #, normalize=True, range=(0,1))
+
+                elif args.save_option == "image_only":
+                    vutils.save_image(g_img, 
+                        os.path.join(dist, '%d_img.png'%(i*args.batch+j)))#, normalize=True, range=(-1,1))
+                    
+                elif args.save_option == "mask_only":
+                    vutils.save_image(g_mask, 
+                        os.path.join(dist, '%d_mask.png'%(i*args.batch+j)))#, normalize=True, range=(-1,1))
+                else:
+                    print("wrong choise to save option.")