chore(inference) : add gradio model

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model/*.model filter=lfs diff=lfs merge=lfs -text
+model/*.pth filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,44 @@
+import torch, os, gradio as gr, numpy as np
+from torchvision import utils, transforms
+from progan_modules import Generator
+
+CHECKPOINT_DIR = "./model"
+Z_DIM, CHANNEL_SIZE = 128, 128
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+FIXED_STEP  = 6
+FIXED_ALPHA = 0.0
+
+g_running = Generator(CHANNEL_SIZE, Z_DIM, pixel_norm=False, tanh=False).to(DEVICE)
+g_running.load_state_dict(torch.load(os.path.join(CHECKPOINT_DIR, "g.model"), map_location=DEVICE))
+g_running.eval()
+
+to_pil = transforms.ToPILImage()
+
+@torch.inference_mode()
+def sample_images(n_images: int = 50, seed: int | None = None):
+    if seed is not None and seed >= 0:
+        torch.manual_seed(seed); np.random.seed(seed)
+    else:
+        torch.seed()
+
+    z = torch.randn(n_images, Z_DIM, device=DEVICE)
+    imgs = g_running(z, step=FIXED_STEP, alpha=FIXED_ALPHA).cpu()
+
+    grid = utils.make_grid(imgs, nrow=10, normalize=True, value_range=(-1, 1))
+    return to_pil(grid)
+
+demo = gr.Interface(
+    fn=sample_images,
+    inputs=[
+        gr.Slider(1, 200, value=50, step=10, label="Jumlah Gambar (kelipatan 10)"),
+        gr.Number(value=-1, precision=0, label="Seed (‑1 = acak)"),
+    ],
+    outputs=gr.Image(type="pil", label="Grid Hasil"),
+    title="Progressive Growing Generative Adversarial Network",
+    description="contoh implementasi PGGAN untuk dataset jerawat",
+    allow_flagging="never",
+)
+
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch(show_api=False, share=True)

model/d.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1ff8ae7d55d9126ccf99e1177d9e63f6884ab9404dc9501ff62b5d5752628cd
+size 6396418

model/d_optim.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3165c039f10cd0d541e88b3dfbf6dfe2a13e0aca6e7f771d6eaa1f610f451b04
+size 12640648

model/g.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e14508dc41768a63d4a5021547215212ba708df99efe53cee3bbeefbe54e188b
+size 6422598

model/g_optim.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92601d2b20eaade8f3f94aea3077019fdafecd59799095ff13ed5236a9308c54
+size 12694058

progan_modules.py

@@ -0,0 +1,250 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from math import sqrt
+
+
+class EqualLR:
+    def __init__(self, name):
+        self.name = name
+
+    def compute_weight(self, module):
+        weight = getattr(module, self.name + '_orig')
+        fan_in = weight.data.size(1) * weight.data[0][0].numel()
+
+        return weight * sqrt(2 / fan_in)
+
+    @staticmethod
+    def apply(module, name):
+        fn = EqualLR(name)
+
+        weight = getattr(module, name)
+        del module._parameters[name]
+        module.register_parameter(name + '_orig', nn.Parameter(weight.data))
+        module.register_forward_pre_hook(fn)
+
+        return fn
+
+    def __call__(self, module, input):
+        weight = self.compute_weight(module)
+        setattr(module, self.name, weight)
+
+
+def equal_lr(module, name='weight'):
+    EqualLR.apply(module, name)
+
+    return module
+
+
+class PixelNorm(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, input):
+        return input / torch.sqrt(torch.mean(input ** 2, dim=1, keepdim=True)
+                                  + 1e-8)
+
+
+class EqualConv2d(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+        conv = nn.Conv2d(*args, **kwargs)
+        conv.weight.data.normal_()
+        conv.bias.data.zero_()
+        self.conv = equal_lr(conv)
+
+    def forward(self, input):
+        return self.conv(input)
+
+
+class EqualConvTranspose2d(nn.Module):
+    ### additional module for OOGAN usage
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+        conv = nn.ConvTranspose2d(*args, **kwargs)
+        conv.weight.data.normal_()
+        conv.bias.data.zero_()
+        self.conv = equal_lr(conv)
+
+    def forward(self, input):
+        return self.conv(input)
+
+class EqualLinear(nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+
+        linear = nn.Linear(in_dim, out_dim)
+        linear.weight.data.normal_()
+        linear.bias.data.zero_()
+
+        self.linear = equal_lr(linear)
+
+    def forward(self, input):
+        return self.linear(input)
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, kernel_size, padding, kernel_size2=None, padding2=None, pixel_norm=True):
+        super().__init__()
+
+        pad1 = padding
+        pad2 = padding
+        if padding2 is not None:
+            pad2 = padding2
+
+        kernel1 = kernel_size
+        kernel2 = kernel_size
+        if kernel_size2 is not None:
+            kernel2 = kernel_size2
+
+        convs = [EqualConv2d(in_channel, out_channel, kernel1, padding=pad1)]
+        if pixel_norm:
+            convs.append(PixelNorm())
+        convs.append(nn.LeakyReLU(0.1))
+        convs.append(EqualConv2d(out_channel, out_channel, kernel2, padding=pad2))
+        if pixel_norm:
+            convs.append(PixelNorm())
+        convs.append(nn.LeakyReLU(0.1))
+
+        self.conv = nn.Sequential(*convs)
+
+    def forward(self, input):
+        out = self.conv(input)
+        return out
+
+
+def upscale(feat):
+    return F.interpolate(feat, scale_factor=2, mode='bilinear', align_corners=False)
+
+class Generator(nn.Module):
+    def __init__(self, input_code_dim=128, in_channel=128, pixel_norm=True, tanh=True):
+        super().__init__()
+        self.input_dim = input_code_dim
+        self.tanh = tanh
+        self.input_layer = nn.Sequential(
+            EqualConvTranspose2d(input_code_dim, in_channel, 4, 1, 0),
+            PixelNorm(),
+            nn.LeakyReLU(0.1))
+
+        self.progression_4 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
+        self.progression_8 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
+        self.progression_16 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
+        self.progression_32 = ConvBlock(in_channel, in_channel, 3, 1, pixel_norm=pixel_norm)
+        self.progression_64 = ConvBlock(in_channel, in_channel//2, 3, 1, pixel_norm=pixel_norm)
+        self.progression_128 = ConvBlock(in_channel//2, in_channel//4, 3, 1, pixel_norm=pixel_norm)
+        self.progression_256 = ConvBlock(in_channel//4, in_channel//4, 3, 1, pixel_norm=pixel_norm)
+
+        self.to_rgb_8 = EqualConv2d(in_channel, 3, 1)
+        self.to_rgb_16 = EqualConv2d(in_channel, 3, 1)
+        self.to_rgb_32 = EqualConv2d(in_channel, 3, 1)
+        self.to_rgb_64 = EqualConv2d(in_channel//2, 3, 1)
+        self.to_rgb_128 = EqualConv2d(in_channel//4, 3, 1)
+        self.to_rgb_256 = EqualConv2d(in_channel//4, 3, 1)
+        
+        self.max_step = 6
+
+    def progress(self, feat, module):
+        out = F.interpolate(feat, scale_factor=2, mode='bilinear', align_corners=False)
+        out = module(out)
+        return out
+
+    def output(self, feat1, feat2, module1, module2, alpha):
+        if 0 <= alpha < 1:
+            skip_rgb = upscale(module1(feat1))
+            out = (1-alpha)*skip_rgb + alpha*module2(feat2)
+        else:
+            out = module2(feat2)
+        if self.tanh:
+            return torch.tanh(out)
+        return out
+
+    def forward(self, input, step=0, alpha=-1):
+        if step > self.max_step:
+            step = self.max_step
+
+        out_4 = self.input_layer(input.view(-1, self.input_dim, 1, 1))
+        out_4 = self.progression_4(out_4)
+        out_8 = self.progress(out_4, self.progression_8)
+        if step==1:
+            if self.tanh:
+                return torch.tanh(self.to_rgb_8(out_8))
+            return self.to_rgb_8(out_8)
+        
+        out_16 = self.progress(out_8, self.progression_16)
+        if step==2:
+            return self.output( out_8, out_16, self.to_rgb_8, self.to_rgb_16, alpha )
+        
+        out_32 = self.progress(out_16, self.progression_32)
+        if step==3:
+            return self.output( out_16, out_32, self.to_rgb_16, self.to_rgb_32, alpha )
+
+        out_64 = self.progress(out_32, self.progression_64)
+        if step==4:
+            return self.output( out_32, out_64, self.to_rgb_32, self.to_rgb_64, alpha )
+        
+        out_128 = self.progress(out_64, self.progression_128)
+        if step==5:
+            return self.output( out_64, out_128, self.to_rgb_64, self.to_rgb_128, alpha )
+
+        out_256 = self.progress(out_128, self.progression_256)
+        if step==6:
+            return self.output( out_128, out_256, self.to_rgb_128, self.to_rgb_256, alpha )
+
+
+class Discriminator(nn.Module):
+    def __init__(self, feat_dim=128):
+        super().__init__()
+
+        self.progression = nn.ModuleList([ConvBlock(feat_dim//4, feat_dim//4, 3, 1),
+                                          ConvBlock(feat_dim//4, feat_dim//2, 3, 1),
+                                          ConvBlock(feat_dim//2, feat_dim, 3, 1),
+                                          ConvBlock(feat_dim, feat_dim, 3, 1),
+                                          ConvBlock(feat_dim, feat_dim, 3, 1),
+                                          ConvBlock(feat_dim, feat_dim, 3, 1),
+                                          ConvBlock(feat_dim+1, feat_dim, 3, 1, 4, 0)])
+
+        self.from_rgb = nn.ModuleList([EqualConv2d(3, feat_dim//4, 1),
+                                       EqualConv2d(3, feat_dim//4, 1),
+                                       EqualConv2d(3, feat_dim//2, 1),
+                                       EqualConv2d(3, feat_dim, 1),
+                                       EqualConv2d(3, feat_dim, 1),
+                                       EqualConv2d(3, feat_dim, 1),
+                                       EqualConv2d(3, feat_dim, 1)])
+
+        self.n_layer = len(self.progression)
+
+        self.linear = EqualLinear(feat_dim, 1)
+
+    def forward(self, input, step=0, alpha=-1):
+        for i in range(step, -1, -1):
+            index = self.n_layer - i - 1
+
+            if i == step:
+                out = self.from_rgb[index](input)
+
+            if i == 0:
+                out_std = torch.sqrt(out.var(0, unbiased=False) + 1e-8)
+                mean_std = out_std.mean()
+                mean_std = mean_std.expand(out.size(0), 1, 4, 4)
+                out = torch.cat([out, mean_std], 1)
+
+            out = self.progression[index](out)
+
+            if i > 0:
+                # out = F.avg_pool2d(out, 2)
+                out = F.interpolate(out, scale_factor=0.5, mode='bilinear', align_corners=False)
+
+                if i == step and 0 <= alpha < 1:
+                    # skip_rgb = F.avg_pool2d(input, 2)
+                    skip_rgb = F.interpolate(input, scale_factor=0.5, mode='bilinear', align_corners=False)
+                    skip_rgb = self.from_rgb[index + 1](skip_rgb)
+                    out = (1 - alpha) * skip_rgb + alpha * out
+
+        out = out.squeeze(2).squeeze(2)
+        # print(input.size(), out.size(), step)
+        out = self.linear(out)
+
+        return out

requirements.txt

@@ -0,0 +1,8 @@
+torch==2.0.0
+torchvision==0.15.0
+numpy==1.23.4
+Pillow==9.2.0
+tqdm==4.64.1
+pydantic==2.10.6
+gradio==3.32.0
+gradio_client>=0.13.0

train.py

@@ -0,0 +1,281 @@
+from tqdm import tqdm
+import numpy as np
+from PIL import Image
+import argparse
+import random
+import torch
+import torch.nn.functional as F
+import os
+from torch import nn, optim
+from torch.autograd import Variable, grad
+from torch.utils.data import DataLoader
+from torchvision import datasets, transforms, utils
+
+from progan_modules import Generator, Discriminator
+
+
+def accumulate(model1, model2, decay=0.999):
+    par1 = dict(model1.named_parameters())
+    par2 = dict(model2.named_parameters())
+
+    for k in par1.keys():
+        par1[k].data.mul_(decay).add_(par2[k].data, alpha=(1 - decay))
+
+
+def imagefolder_loader(path):
+    def loader(transform):
+        data = datasets.ImageFolder(path, transform=transform)
+        data_loader = DataLoader(data, shuffle=True, batch_size=batch_size, num_workers=2)
+        return data_loader
+    return loader
+
+
+def sample_data(dataloader, image_size=4):
+    transform = transforms.Compose([
+        transforms.Resize(image_size+int(image_size*0.2)+1),
+        transforms.RandomCrop(image_size),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
+    ])
+
+    loader = dataloader(transform)
+
+    return loader
+
+
+def train(generator, discriminator, init_step, loader, total_iter=600000, start_iter=0):
+    step = init_step # can be 1 = 8, 2 = 16, 3 = 32, 4 = 64, 5 = 128, 6 = 128
+    data_loader = sample_data(loader, 4 * 2 ** step)
+    dataset = iter(data_loader)
+
+    #total_iter = 600000
+    total_iter_remain = total_iter - (total_iter//6)*(step-1)
+
+    pbar = tqdm(range(total_iter_remain))
+
+    disc_loss_val = 0
+    gen_loss_val = 0
+    grad_loss_val = 0
+
+    from datetime import datetime
+    import os
+    date_time = datetime.now()
+    post_fix = '%s_%s_%d_%d.txt'%(trial_name, date_time.date(), date_time.hour, date_time.minute)
+    log_folder = 'trial_%s_%s_%d_%d'%(trial_name, date_time.date(), date_time.hour, date_time.minute)
+    
+    os.mkdir(log_folder)
+    os.mkdir(log_folder+'/checkpoint')
+    os.mkdir(log_folder+'/sample')
+
+    config_file_name = os.path.join(log_folder, 'train_config_'+post_fix)
+    config_file = open(config_file_name, 'w')
+    config_file.write(str(args))
+    config_file.close()
+
+    log_file_name = os.path.join(log_folder, 'train_log_'+post_fix)
+    log_file = open(log_file_name, 'w')
+    log_file.write('g,d,nll,onehot\n')
+    log_file.close()
+
+    from shutil import copy
+    copy('train.py', log_folder+'/train_%s.py'%post_fix)
+    copy('progan_modules.py', log_folder+'/model_%s.py'%post_fix)
+
+    alpha = 0
+    #one = torch.FloatTensor([1]).to(device)
+    one = torch.tensor(1, dtype=torch.float).to(device)
+    mone = one * -1
+    iteration = 0
+
+    for i in pbar:
+        discriminator.zero_grad()
+
+        alpha = min(1, (2/(total_iter//6)) * iteration)
+
+        if iteration > total_iter//6:
+            alpha = 0
+            iteration = 0
+            step += 1
+
+            if step > 6:
+                alpha = 1
+                step = 6
+            data_loader = sample_data(loader, 4 * 2 ** step)
+            dataset = iter(data_loader)
+
+        try:
+            real_image, label = next(dataset)
+
+        except (OSError, StopIteration):
+            dataset = iter(data_loader)
+            real_image, label = next(dataset)
+
+        iteration += 1
+
+        ### 1. train Discriminator
+        b_size = real_image.size(0)
+        real_image = real_image.to(device)
+        label = label.to(device)
+        real_predict = discriminator(
+            real_image, step=step, alpha=alpha)
+        real_predict = real_predict.mean() \
+            - 0.001 * (real_predict ** 2).mean()
+        real_predict.backward(mone)
+
+        # sample input data: vector for Generator
+        gen_z = torch.randn(b_size, input_code_size).to(device)
+
+        fake_image = generator(gen_z, step=step, alpha=alpha)
+        fake_predict = discriminator(
+            fake_image.detach(), step=step, alpha=alpha)
+        fake_predict = fake_predict.mean()
+        fake_predict.backward(one)
+
+        ### gradient penalty for D
+        eps = torch.rand(b_size, 1, 1, 1).to(device)
+        x_hat = eps * real_image.data + (1 - eps) * fake_image.detach().data
+        x_hat.requires_grad = True
+        hat_predict = discriminator(x_hat, step=step, alpha=alpha)
+        grad_x_hat = grad(
+            outputs=hat_predict.sum(), inputs=x_hat, create_graph=True)[0]
+        grad_penalty = ((grad_x_hat.view(grad_x_hat.size(0), -1)
+                         .norm(2, dim=1) - 1)**2).mean()
+        grad_penalty = 10 * grad_penalty
+        grad_penalty.backward()
+        grad_loss_val += grad_penalty.item()
+        disc_loss_val += (real_predict - fake_predict).item()
+
+        d_optimizer.step()
+
+        ### 2. train Generator
+        if (i + 1) % n_critic == 0:
+            generator.zero_grad()
+            discriminator.zero_grad()
+            
+            predict = discriminator(fake_image, step=step, alpha=alpha)
+
+            loss = -predict.mean()
+            gen_loss_val += loss.item()
+
+
+            loss.backward()
+            g_optimizer.step()
+            accumulate(g_running, generator)
+
+        if (i + 1) % 1000 == 0 or i==0:
+            with torch.no_grad():
+                images = g_running(torch.randn(5 * 10, input_code_size).to(device), step=step, alpha=alpha).data.cpu()
+
+                utils.save_image(
+                    images,
+                    f'{log_folder}/sample/{str(i + 1).zfill(6)}.png',
+                    nrow=10,
+                    normalize=True)
+ 
+        if (i+1) % 10000 == 0 or i==0:
+            try:
+                torch.save(g_running.state_dict(), f'{log_folder}/checkpoint/{str(i + 1).zfill(6)}_g.model')
+                torch.save(discriminator.state_dict(), f'{log_folder}/checkpoint/{str(i + 1).zfill(6)}_d.model')
+                torch.save(g_optimizer.state_dict(), os.path.join(log_folder, 'checkpoint', f'{str(i + 1).zfill(6)}_g_optim.pth'))
+                torch.save(d_optimizer.state_dict(), os.path.join(log_folder, 'checkpoint', f'{str(i + 1).zfill(6)}_d_optim.pth'))
+            except:
+                pass
+
+        if (i+1)%500 == 0:
+            state_msg = (f'{i + 1}; G: {gen_loss_val/(500//n_critic):.3f}; D: {disc_loss_val/500:.3f};'
+                f' Grad: {grad_loss_val/500:.3f}; Alpha: {alpha:.3f}')
+            
+            log_file = open(log_file_name, 'a+')
+            new_line = "%.5f,%.5f\n"%(gen_loss_val/(500//n_critic), disc_loss_val/500)
+            log_file.write(new_line)
+            log_file.close()
+
+            disc_loss_val = 0
+            gen_loss_val = 0
+            grad_loss_val = 0
+
+            print(state_msg)
+            #pbar.set_description(state_msg)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Progressive GAN, during training, the model will learn to generate  images from a low resolution, then progressively getting high resolution ')
+
+    parser.add_argument('--start_iter', type=int, default=0, help='Iterasi awal dari training')
+    parser.add_argument('--checkpoint', type=str, default="/content/model/", help='Path to model checkpoint directory (default: None, train from scratch)')
+    parser.add_argument('--path', type=str,default="/content/merged_dataset/Acne", help='path of specified dataset, should be a folder that has one or many sub image folders inside')
+    parser.add_argument('--trial_name', type=str, default="test1", help='a brief description of the training trial')
+    parser.add_argument('--gpu_id', type=int, default=0, help='0 is the first gpu, 1 is the second gpu, etc.')
+    parser.add_argument('--lr', type=float, default=0.001, help='learning rate, default is 1e-3, usually dont need to change it, you can try make it bigger, such as 2e-3')
+    parser.add_argument('--z_dim', type=int, default=128, help='the initial latent vector\'s dimension, can be smaller such as 64, if the dataset is not diverse')
+    parser.add_argument('--channel', type=int, default=128, help='determines how big the model is, smaller value means faster training, but less capacity of the model')
+    parser.add_argument('--batch_size', type=int, default=4, help='how many images to train together at one iteration')
+    parser.add_argument('--n_critic', type=int, default=1, help='train Dhow many times while train G 1 time')
+    parser.add_argument('--init_step', type=int, default=1, help='start from what resolution, 1 means 8x8 resolution, 2 means 16x16 resolution, ..., 6 means 256x256 resolution')
+    parser.add_argument('--total_iter', type=int, default=300000, help='how many iterations to train in total, the value is in assumption that init step is 1')
+    parser.add_argument('--pixel_norm', default=False, action="store_true", help='a normalization method inside the model, you can try use it or not depends on the dataset')
+    parser.add_argument('--tanh', default=False, action="store_true", help='an output non-linearity on the output of Generator, you can try use it or not depends on the dataset')
+    
+    args = parser.parse_args()
+
+    trial_name = args.trial_name
+    device = torch.device("cuda:%d"%(args.gpu_id))
+    input_code_size = args.z_dim
+    batch_size = args.batch_size
+    n_critic = args.n_critic
+
+    generator = Generator(in_channel=args.channel, input_code_dim=input_code_size, pixel_norm=args.pixel_norm, tanh=args.tanh).to(device)
+    discriminator = Discriminator(feat_dim=args.channel).to(device)
+    g_running = Generator(in_channel=args.channel, input_code_dim=input_code_size, pixel_norm=args.pixel_norm, tanh=args.tanh).to(device)
+    
+    ## you can directly load a pretrained model here
+    if args.checkpoint:
+        generator_path = os.path.join(args.checkpoint, "g.model")
+        discriminator_path = os.path.join(args.checkpoint, "d.model")
+
+        if os.path.exists(generator_path) and os.path.exists(discriminator_path):
+            print(f"Loading checkpoints from {args.checkpoint}...")
+            generator.load_state_dict(torch.load(generator_path))
+            g_running.load_state_dict(torch.load(generator_path))
+            discriminator.load_state_dict(torch.load(discriminator_path))
+        else:
+            print(f"Warning: Checkpoint not found at {args.checkpoint}. Training from scratch!")
+    else:
+        print("No checkpoint provided, training from scratch.")
+    
+    if args.checkpoint:
+        generator_path = os.path.join(args.checkpoint, "g.model")
+        discriminator_path = os.path.join(args.checkpoint, "d.model")
+        optimizer_g_path = os.path.join(args.checkpoint, "g_optim.pth")
+        optimizer_d_path = os.path.join(args.checkpoint, "d_optim.pth")
+
+        if os.path.exists(generator_path) and os.path.exists(discriminator_path):
+            print(f"Loading checkpoints from {args.checkpoint}...")
+            generator.load_state_dict(torch.load(generator_path))
+            g_running.load_state_dict(torch.load(generator_path))
+            discriminator.load_state_dict(torch.load(discriminator_path))
+        else:
+            print(f"Warning: Checkpoint not found at {args.checkpoint}. Training from scratch!")
+    else:
+        print("No checkpoint provided, training from scratch.")
+
+    g_running.train(False)
+
+    g_optimizer = optim.Adam(generator.parameters(), lr=args.lr, betas=(0.0, 0.99))
+    d_optimizer = optim.Adam(discriminator.parameters(), lr=args.lr, betas=(0.0, 0.99))
+
+    optimizer_g_path = os.path.join(args.checkpoint, "g_optim.pth")
+    optimizer_d_path = os.path.join(args.checkpoint, "d_optim.pth")
+    
+    if os.path.exists(optimizer_g_path) and os.path.exists(optimizer_d_path):
+        g_optimizer.load_state_dict(torch.load(optimizer_g_path))
+        d_optimizer.load_state_dict(torch.load(optimizer_d_path))
+        print("Optimizers loaded successfully!")
+    else:
+        print("Warning: Optimizer checkpoint not found. Using new optimizers!")
+    accumulate(g_running, generator, 0)
+
+    loader = imagefolder_loader(args.path)
+
+    train(generator, discriminator, args.init_step, loader, args.total_iter, args.start_iter)