libcom 0.1.8__py3-none-any.whl

libcom/painterly_image_harmonization/source/PHDNet/networks.py

@@ -0,0 +1,202 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def calc_mean_std(feat, eps=1e-5):
+    # eps is a small value added to the variance to avoid divide-by-zero.
+    size = feat.size()
+    assert (len(size) == 4)
+    N, C = size[:2]
+    feat_var = feat.contiguous().view(N, C, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().view(N, C, 1, 1)
+    feat_mean = feat.contiguous().view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
+    return feat_mean, feat_std
+
+def adain_fg(comp_feat, mask):
+    #assert (content_feat.size()[:2] == style_feat.size()[:2])
+    size = comp_feat.size()
+    #style_mean, style_std = calc_mean_std(style_feat)  # the style features
+    downsample_mask_style = 1 - mask
+    style_mean, style_std = get_foreground_mean_std(comp_feat, downsample_mask_style)  # the style features
+    fore_mean, fore_std = get_foreground_mean_std(comp_feat, mask)  # the foreground features
+
+    normalized_feat = (comp_feat - fore_mean.expand(size)) / fore_std.expand(size)
+    return (normalized_feat * style_std.expand(size) + style_mean.expand(size)) * mask + (comp_feat * (1 - mask))
+
+
+decoder_cat = nn.Sequential(
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 256, (3, 3)),
+    nn.ReLU(), # relu1-1
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),  # 4
+    nn.Conv2d(512, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 128, (3, 3)),
+    nn.ReLU(), # relu2-1
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),  # 17
+    nn.Conv2d(256, 128, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 64, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),  # 24
+    nn.Conv2d(128, 64, (3, 3)),
+    nn.ReLU(),
+    nn.Conv2d(65, 1, (1, 1),padding=0,stride=1), ##matting layer
+    nn.ReflectionPad2d((1, 1, 1, 1)), # 29
+    nn.Conv2d(64, 3, (3, 3)),
+)
+
+
+vgg = nn.Sequential(
+    nn.Conv2d(3, 3, (1, 1)),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(3, 64, (3, 3)),
+    nn.ReLU(),  # relu1-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 64, (3, 3)),
+    nn.ReLU(),  # relu1-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 128, (3, 3)),
+    nn.ReLU(),  # relu2-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 128, (3, 3)),
+    nn.ReLU(),  # relu2-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 256, (3, 3)),
+    nn.ReLU(),  # relu3-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)), 
+    nn.ReLU(),  # relu3-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 512, (3, 3)),
+    nn.ReLU(),  # relu4-1, this is the last layer used
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU()  # relu5-4
+)
+
+
+
+def get_foreground_mean_std(features, mask, eps=1e-5):
+    region = features * mask 
+    sum = torch.sum(region, dim=[2, 3])     # (B, C)
+    num = torch.sum(mask, dim=[2, 3])       # (B, C)
+    mu = sum / (num + eps)
+    mean = mu[:, :, None, None]
+    var = torch.sum((region + (1 - mask)*mean - mean) ** 2, dim=[2, 3]) / (num + eps)
+    var = var[:, :, None, None]
+    std = torch.sqrt(var+eps)
+    return mean, std
+
+
+class PHDNet(nn.Module):
+    def __init__(self, vgg, decoder):
+        super(PHDNet, self).__init__()
+        # load the pretrained VGG encoder
+        vgg_layers = list(vgg.children())
+        self.enc_1 = nn.Sequential(*vgg_layers[:4])  # input -> relu1_1
+        self.enc_2 = nn.Sequential(*vgg_layers[4:11])  # relu1_1 -> relu2_1
+        self.enc_3 = nn.Sequential(*vgg_layers[11:18])  # relu2_1 -> relu3_1
+        self.enc_4 = nn.Sequential(*vgg_layers[18:31])  # relu3_1 -> relu4_1
+       
+
+        # define the decoder
+        self.decoder = decoder
+        dec_layers = list(decoder.children())
+        self.dec_1 =  nn.Sequential(*dec_layers[:4]) 
+        self.dec_2 =  nn.Sequential(*dec_layers[4:17]) 
+        self.dec_3 =  nn.Sequential(*dec_layers[17:24]) 
+        self.dec_4 =  nn.Sequential(*dec_layers[24:27]) 
+        self.conv_attention = nn.Sequential(*dec_layers[27:28]) 
+        self.dec_4_2 =  nn.Sequential(*dec_layers[28:]) 
+
+        # fix the VGG encoder
+        for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4']:
+            for param in getattr(self, name).parameters():
+                param.requires_grad = False
+
+
+    # extract relu1_1, relu2_1, relu3_1, relu4_1 from input image
+    def encode_with_intermediate(self, input):
+        results = [input]
+        for i in range(4):
+            func = getattr(self, 'enc_{:d}'.format(i + 1))
+            results.append(func(results[-1]))
+        return results[1:]
+    
+    
+    def decode(self, comp, mask, enc_features):
+        width = height = enc_features[-1].size(-1)
+        downsample_mask = self.downsample(mask, width, height)
+        t = adain_fg(enc_features[-1], downsample_mask)
+        dec_feature = self.dec_1(t)
+
+        for i in range(1,4):
+            func = getattr(self, 'dec_{:d}'.format(i + 1))
+            width = height = enc_features[-(i+1)].size(-1)
+            downsample_mask = self.downsample(mask, width, height)
+            t = adain_fg(enc_features[-(i+1)], downsample_mask)
+            dec_feature = func(torch.cat([dec_feature, t], dim=1))
+
+        
+        width = height = dec_feature.size(-1)
+        downsample_mask = self.downsample(mask, width, height)
+
+        attention_map = torch.sigmoid(self.conv_attention(torch.cat([dec_feature, downsample_mask], dim=1)))
+        coarse_output = self.dec_4_2(dec_feature)
+        output = attention_map * coarse_output + (1.0 - attention_map) * comp
+        return output
+
+
+    def downsample(self, image_tensor, width, height):
+        image_upsample_tensor = torch.nn.functional.interpolate(image_tensor, size=[width, height])
+        image_upsample_tensor = image_upsample_tensor.clamp(0, 1)
+        return image_upsample_tensor
+
+
+    def forward(self, comp, mask):
+        comb_feats = self.encode_with_intermediate(comp)
+        final_output = self.decode(comp, mask, comb_feats)
+
+        return final_output
+

libcom/painterly_image_harmonization/source/PHDNet/phdnet.py

@@ -0,0 +1,47 @@
+import torch
+from . import networks
+from torch import nn
+import os
+
+class PHDNet(nn.Module):
+    def __init__(self, device):
+        super(PHDNet, self).__init__()
+        # define networks
+        self.device = device
+        self.netvgg = networks.vgg
+        self.netvgg = nn.Sequential(*list(self.netvgg.children())[:31])
+        self.netDecoder = networks.decoder_cat
+        self.netG = networks.PHDNet(self.netvgg, self.netDecoder)
+
+
+    def forward(self, comp, mask):
+        """Employ generator to generate the output"""
+        if len(comp.shape) == 3:
+            comp = comp.unsqueeze(0)
+            mask = mask.unsqueeze(0)
+
+        self.comp = comp.to(self.device)
+        self.mask = mask.to(self.device)
+        
+        output = self.netG(self.comp, self.mask)
+        return output
+
+
+    def load_networks(self, load_path):
+        """Load all the networks from the disk.
+
+        Parameters:
+            load_path (str) -- path to load parameters
+        """
+        if os.path.exists(load_path):
+            net = getattr(self, 'netG')
+            if isinstance(net, torch.nn.DataParallel):
+                net = net.module
+            # print('loading the model from %s' % load_path)
+            state_dict = torch.load(load_path, map_location=self.device,weights_only=True)
+            if hasattr(state_dict, '_metadata'):
+                del state_dict._metadata
+
+            net.load_state_dict(state_dict, strict=True)
+        else:
+            print('failed loading model from {}'.format(load_path))

libcom/painterly_image_harmonization/source/PHDNet/phdnet_config.py

@@ -0,0 +1,18 @@
+import os
+
+class Config(object):
+    # config
+    backbone = 'resnet18'
+    model = 'phdnet'
+    dataset_mode = 'phd'
+    preprocess = 'none'
+    batch_size = 1
+    num_threads = 6
+    gpu_ids = [2]
+    load_size = 512
+
+
+opt = Config()
+
+if __name__ == "__main__":
+    print(opt.exp_root)

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/inference_base.py

@@ -0,0 +1,300 @@
+import argparse
+import torch
+from omegaconf import OmegaConf
+
+from libcom.painterly_image_harmonization.source.ldm.models.diffusion.ddim import DDIMSampler
+from libcom.painterly_image_harmonization.source.ldm.models.diffusion.plms import PLMSSampler
+from libcom.painterly_image_harmonization.source.ldm.modules.encoders.adapter import Adapter, StyleAdapter, Adapter_light
+from libcom.painterly_image_harmonization.source.ldm.modules.extra_condition.api import ExtraCondition
+from libcom.painterly_image_harmonization.source.ldm.util import fix_cond_shapes, load_model_from_config, read_state_dict
+
+DEFAULT_NEGATIVE_PROMPT = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, ' \
+                          'fewer digits, cropped, worst quality, low quality'
+
+
+def get_base_argument_parser() -> argparse.ArgumentParser:
+    """get the base argument parser for inference scripts"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--outdir',
+        type=str,
+        help='dir to write results to',
+        default=None,
+    )
+
+    parser.add_argument(
+        '--prompt',
+        type=str,
+        nargs='?',
+        default=None,
+        help='positive prompt',
+    )
+
+    parser.add_argument(
+        '--neg_prompt',
+        type=str,
+        default=DEFAULT_NEGATIVE_PROMPT,
+        help='negative prompt',
+    )
+
+    parser.add_argument(
+        '--cond_path',
+        type=str,
+        default=None,
+        help='condition image path',
+    )
+
+    parser.add_argument(
+        '--cond_inp_type',
+        type=str,
+        default='image',
+        help='the type of the input condition image, take depth T2I as example, the input can be raw image, '
+        'which depth will be calculated, or the input can be a directly a depth map image',
+    )
+
+    parser.add_argument(
+        '--sampler',
+        type=str,
+        default='ddim',
+        choices=['ddim', 'plms'],
+        help='sampling algorithm, currently, only ddim and plms are supported, more are on the way',
+    )
+
+    parser.add_argument(
+        '--steps',
+        type=int,
+        default=50,
+        help='number of sampling steps',
+    )
+
+    parser.add_argument(
+        '--sd_ckpt',
+        type=str,
+        default='models/sd-v1-4.ckpt',
+        help='path to checkpoint of stable diffusion model, both .ckpt and .safetensor are supported',
+    )
+
+    parser.add_argument(
+        '--vae_ckpt',
+        type=str,
+        default=None,
+        help='vae checkpoint, anime SD models usually have seperate vae ckpt that need to be loaded',
+    )
+
+    parser.add_argument(
+        '--adapter_ckpt',
+        type=str,
+        default=None,
+        help='path to checkpoint of adapter',
+    )
+
+    parser.add_argument(
+        '--config',
+        type=str,
+        default='configs/stable-diffusion/sd-v1-inference.yaml',
+        help='path to config which constructs SD model',
+    )
+
+    parser.add_argument(
+        '--max_resolution',
+        type=float,
+        default=512 * 512,
+        help='max image height * width, only for computer with limited vram',
+    )
+
+    parser.add_argument(
+        '--resize_short_edge',
+        type=int,
+        default=None,
+        help='resize short edge of the input image, if this arg is set, max_resolution will not be used',
+    )
+
+    parser.add_argument(
+        '--C',
+        type=int,
+        default=4,
+        help='latent channels',
+    )
+
+    parser.add_argument(
+        '--f',
+        type=int,
+        default=8,
+        help='downsampling factor',
+    )
+
+    parser.add_argument(
+        '--scale',
+        type=float,
+        default=7.5,
+        help='unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))',
+    )
+
+    parser.add_argument(
+        '--cond_tau',
+        type=float,
+        default=1.0,
+        help='timestamp parameter that determines until which step the adapter is applied, '
+        'similar as Prompt-to-Prompt tau',
+    )
+
+    parser.add_argument(
+        '--style_cond_tau',
+        type=float,
+        default=1.0,
+        help='timestamp parameter that determines until which step the adapter is applied, '
+             'similar as Prompt-to-Prompt tau',
+    )
+
+    parser.add_argument(
+        '--cond_weight',
+        type=float,
+        default=1.0,
+        help='the adapter features are multiplied by the cond_weight. The larger the cond_weight, the more aligned '
+        'the generated image and condition will be, but the generated quality may be reduced',
+    )
+
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=42,
+    )
+
+    parser.add_argument(
+        '--n_samples',
+        type=int,
+        default=4,
+        help='# of samples to generate',
+    )
+
+    return parser
+
+
+def get_sd_models(opt):
+    """
+    build stable diffusion model, sampler
+    """
+    # SD
+    config = OmegaConf.load(f"{opt.config}")
+    model = load_model_from_config(config, opt.sd_ckpt, opt.vae_ckpt)
+    sd_model = model.to(opt.device)
+
+    # sampler
+    if opt.sampler == 'plms':
+        sampler = PLMSSampler(model)
+    elif opt.sampler == 'ddim':
+        sampler = DDIMSampler(model)
+    else:
+        raise NotImplementedError
+
+    return sd_model, sampler
+
+
+def get_t2i_adapter_models(opt):
+    config = OmegaConf.load(f"{opt.config}")
+    model = load_model_from_config(config, opt.sd_ckpt, opt.vae_ckpt)
+    adapter_ckpt_path = getattr(opt, f'{opt.which_cond}_adapter_ckpt', None)
+    if adapter_ckpt_path is None:
+        adapter_ckpt_path = getattr(opt, 'adapter_ckpt')
+    adapter_ckpt = read_state_dict(adapter_ckpt_path)
+    new_state_dict = {}
+    for k, v in adapter_ckpt.items():
+        if not k.startswith('adapter.'):
+            new_state_dict[f'adapter.{k}'] = v
+        else:
+            new_state_dict[k] = v
+    m, u = model.load_state_dict(new_state_dict, strict=False)
+    if len(u) > 0:
+        print(f"unexpected keys in loading adapter ckpt {adapter_ckpt_path}:")
+        print(u)
+
+    model = model.to(opt.device)
+
+    # sampler
+    if opt.sampler == 'plms':
+        sampler = PLMSSampler(model)
+    elif opt.sampler == 'ddim':
+        sampler = DDIMSampler(model)
+    else:
+        raise NotImplementedError
+
+    return model, sampler
+
+
+def get_cond_ch(cond_type: ExtraCondition):
+    if cond_type == ExtraCondition.sketch or cond_type == ExtraCondition.canny:
+        return 1
+    return 3
+
+
+def get_adapters(opt, cond_type: ExtraCondition):
+    adapter = {}
+    cond_weight = getattr(opt, f'{cond_type.name}_weight', None)
+    if cond_weight is None:
+        cond_weight = getattr(opt, 'cond_weight')
+    adapter['cond_weight'] = cond_weight
+
+    if cond_type == ExtraCondition.style:
+        adapter['model'] = StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8).to(opt.device)
+    elif cond_type == ExtraCondition.color:
+        adapter['model'] = Adapter_light(
+            cin=64 * get_cond_ch(cond_type),
+            channels=[320, 640, 1280, 1280],
+            nums_rb=4).to(opt.device)
+    else:
+        adapter['model'] = Adapter(
+            cin=64 * get_cond_ch(cond_type),
+            channels=[320, 640, 1280, 1280][:4],
+            nums_rb=2,
+            ksize=1,
+            sk=True,
+            use_conv=False).to(opt.device)
+    ckpt_path = getattr(opt, f'{cond_type.name}_adapter_ckpt', None)
+    if ckpt_path is None:
+        ckpt_path = getattr(opt, 'adapter_ckpt')
+    state_dict = read_state_dict(ckpt_path)
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        if k.startswith('adapter.'):
+            new_state_dict[k[len('adapter.'):]] = v
+        else:
+            new_state_dict[k] = v
+
+    adapter['model'].load_state_dict(new_state_dict)
+
+    return adapter
+
+
+def diffusion_inference(opt, model, sampler, adapter_features, append_to_context=None):
+    # get text embedding
+    c = model.get_learned_conditioning([opt.prompt])
+    if opt.scale != 1.0:
+        uc = model.get_learned_conditioning([opt.neg_prompt])
+    else:
+        uc = None
+    c, uc = fix_cond_shapes(model, c, uc)
+
+    if not hasattr(opt, 'H'):
+        opt.H = 512
+        opt.W = 512
+    shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
+
+    samples_latents, _ = sampler.sample(
+        S=opt.steps,
+        conditioning=c,
+        batch_size=1,
+        shape=shape,
+        verbose=False,
+        unconditional_guidance_scale=opt.scale,
+        unconditional_conditioning=uc,
+        x_T=None,
+        features_adapter=adapter_features,
+        append_to_context=append_to_context,
+        cond_tau=opt.cond_tau,
+        style_cond_tau=opt.style_cond_tau,
+    )
+
+    x_samples = model.decode_first_stage(samples_latents)
+    x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+
+    return x_samples

libcom/painterly_image_harmonization/source/PHDiffusion/ldm/lr_scheduler.py

@@ -0,0 +1,98 @@
+import numpy as np
+
+
+class LambdaWarmUpCosineScheduler:
+    """
+    note: use with a base_lr of 1.0
+    """
+    def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
+        self.lr_warm_up_steps = warm_up_steps
+        self.lr_start = lr_start
+        self.lr_min = lr_min
+        self.lr_max = lr_max
+        self.lr_max_decay_steps = max_decay_steps
+        self.last_lr = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def schedule(self, n, **kwargs):
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
+        if n < self.lr_warm_up_steps:
+            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
+            self.last_lr = lr
+            return lr
+        else:
+            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
+            t = min(t, 1.0)
+            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
+                    1 + np.cos(t * np.pi))
+            self.last_lr = lr
+            return lr
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n,**kwargs)
+
+
+class LambdaWarmUpCosineScheduler2:
+    """
+    supports repeated iterations, configurable via lists
+    note: use with a base_lr of 1.0.
+    """
+    def __init__(self, warm_up_steps, f_min, f_max, f_start, cycle_lengths, verbosity_interval=0):
+        assert len(warm_up_steps) == len(f_min) == len(f_max) == len(f_start) == len(cycle_lengths)
+        self.lr_warm_up_steps = warm_up_steps
+        self.f_start = f_start
+        self.f_min = f_min
+        self.f_max = f_max
+        self.cycle_lengths = cycle_lengths
+        self.cum_cycles = np.cumsum([0] + list(self.cycle_lengths))
+        self.last_f = 0.
+        self.verbosity_interval = verbosity_interval
+
+    def find_in_interval(self, n):
+        interval = 0
+        for cl in self.cum_cycles[1:]:
+            if n <= cl:
+                return interval
+            interval += 1
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                                                       f"current cycle {cycle}")
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            t = (n - self.lr_warm_up_steps[cycle]) / (self.cycle_lengths[cycle] - self.lr_warm_up_steps[cycle])
+            t = min(t, 1.0)
+            f = self.f_min[cycle] + 0.5 * (self.f_max[cycle] - self.f_min[cycle]) * (
+                    1 + np.cos(t * np.pi))
+            self.last_f = f
+            return f
+
+    def __call__(self, n, **kwargs):
+        return self.schedule(n, **kwargs)
+
+
+class LambdaLinearScheduler(LambdaWarmUpCosineScheduler2):
+
+    def schedule(self, n, **kwargs):
+        cycle = self.find_in_interval(n)
+        n = n - self.cum_cycles[cycle]
+        if self.verbosity_interval > 0:
+            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_f}, "
+                                                       f"current cycle {cycle}")
+
+        if n < self.lr_warm_up_steps[cycle]:
+            f = (self.f_max[cycle] - self.f_start[cycle]) / self.lr_warm_up_steps[cycle] * n + self.f_start[cycle]
+            self.last_f = f
+            return f
+        else:
+            f = self.f_min[cycle] + (self.f_max[cycle] - self.f_min[cycle]) * (self.cycle_lengths[cycle] - n) / (self.cycle_lengths[cycle])
+            self.last_f = f
+            return f
+