train/train_ultrapixel_control.py

import torch
import json
import yaml
import torchvision
from torch import nn, optim
from transformers import AutoTokenizer, CLIPTextModelWithProjection, CLIPVisionModelWithProjection
from warmup_scheduler import GradualWarmupScheduler
import torch.multiprocessing as mp
import numpy as np
import sys

import os
from dataclasses import dataclass
from torch.distributed import init_process_group, destroy_process_group, barrier
from gdf import GDF_dual_fixlrt as GDF
from gdf import EpsilonTarget, CosineSchedule
from gdf import VPScaler, CosineTNoiseCond, DDPMSampler, P2LossWeight, AdaptiveLossWeight
from torchtools.transforms import SmartCrop
from fractions import Fraction
from modules.effnet import EfficientNetEncoder

from modules.model_4stage_lite import StageC

from modules.model_4stage_lite import ResBlock, AttnBlock, TimestepBlock, FeedForwardBlock
from modules.common_ckpt import GlobalResponseNorm
from modules.previewer import Previewer
from core.data import Bucketeer
from train.base import DataCore, TrainingCore
from tqdm import tqdm
from core import WarpCore
from core.utils import EXPECTED, EXPECTED_TRAIN, load_or_fail
from torch.distributed.fsdp.wrap import ModuleWrapPolicy, size_based_auto_wrap_policy
from accelerate import init_empty_weights
from accelerate.utils import set_module_tensor_to_device
from contextlib import contextmanager
from train.dist_core import *
import glob
from torch.utils.data import DataLoader, Dataset
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data.distributed import DistributedSampler
from PIL import Image
from core.utils import EXPECTED, EXPECTED_TRAIN, update_weights_ema, create_folder_if_necessary
from core.utils import Base
from modules.common import LayerNorm2d
import torch.nn.functional as F
import functools
import math
import copy
import random
from modules.lora import apply_lora, apply_retoken, LoRA, ReToken
from modules import ControlNet, ControlNetDeliverer
from modules import controlnet_filters

Image.MAX_IMAGE_PIXELS = None
torch.manual_seed(8432)
random.seed(8432)
np.random.seed(8432)
#7978026

class Null_Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x):
        pass


def identity(x):
    if isinstance(x, bytes):
        x = x.decode('utf-8')
    return x
def check_nan_inmodel(model, meta=''):
        for name, param in model.named_parameters():
            if torch.isnan(param).any():
                print(f"nan detected in {name}", meta)
                return True
        print('no nan', meta)
        return False  


class WurstCore(TrainingCore, DataCore, WarpCore):
    @dataclass(frozen=True)
    class Config(TrainingCore.Config, DataCore.Config, WarpCore.Config):
        # TRAINING PARAMS
        lr: float = EXPECTED_TRAIN
        warmup_updates: int = EXPECTED_TRAIN
        dtype: str = None

       # MODEL VERSION
        model_version: str = EXPECTED  # 3.6B or 1B
        clip_image_model_name: str = 'openai/clip-vit-large-patch14'
        clip_text_model_name: str = 'laion/CLIP-ViT-bigG-14-laion2B-39B-b160k'
       
        # CHECKPOINT PATHS
        effnet_checkpoint_path: str = EXPECTED
        previewer_checkpoint_path: str = EXPECTED
        #trans_inr_ckpt: str = EXPECTED
        generator_checkpoint_path: str = None
        controlnet_checkpoint_path: str = EXPECTED
        
        # controlnet settings
        controlnet_blocks: list = EXPECTED
        controlnet_filter: str = EXPECTED
        controlnet_filter_params: dict = None
        controlnet_bottleneck_mode: str = None


        # gdf customization
        adaptive_loss_weight: str = None

        #module_filters: list = EXPECTED
        #rank: int = EXPECTED
    @dataclass(frozen=True)
    class Data(Base):
        dataset: Dataset = EXPECTED
        dataloader: DataLoader  = EXPECTED
        iterator: any = EXPECTED
        sampler: DistributedSampler = EXPECTED

    @dataclass(frozen=True)
    class Models(TrainingCore.Models, DataCore.Models, WarpCore.Models):
        effnet: nn.Module = EXPECTED
        previewer: nn.Module = EXPECTED
        train_norm: nn.Module = EXPECTED
        train_norm_ema: nn.Module = EXPECTED
        controlnet: nn.Module = EXPECTED

    @dataclass(frozen=True)
    class Schedulers(WarpCore.Schedulers):
        generator: any = None

    @dataclass(frozen=True)
    class Extras(TrainingCore.Extras, DataCore.Extras, WarpCore.Extras):
        gdf: GDF = EXPECTED
        sampling_configs: dict = EXPECTED
        effnet_preprocess: torchvision.transforms.Compose = EXPECTED
        controlnet_filter: controlnet_filters.BaseFilter = EXPECTED

    info: TrainingCore.Info
    config: Config

    def setup_extras_pre(self) -> Extras:
        gdf = GDF(
            schedule=CosineSchedule(clamp_range=[0.0001, 0.9999]),
            input_scaler=VPScaler(), target=EpsilonTarget(),
            noise_cond=CosineTNoiseCond(),
            loss_weight=AdaptiveLossWeight() if self.config.adaptive_loss_weight is True else P2LossWeight(),
        )
        sampling_configs = {"cfg": 5, "sampler": DDPMSampler(gdf), "shift": 1, "timesteps": 20}

        if self.info.adaptive_loss is not None:
            gdf.loss_weight.bucket_ranges = torch.tensor(self.info.adaptive_loss['bucket_ranges'])
            gdf.loss_weight.bucket_losses = torch.tensor(self.info.adaptive_loss['bucket_losses'])

        effnet_preprocess = torchvision.transforms.Compose([
            torchvision.transforms.Normalize(
                mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)
            )
        ])

        clip_preprocess = torchvision.transforms.Compose([
            torchvision.transforms.Resize(224, interpolation=torchvision.transforms.InterpolationMode.BICUBIC),
            torchvision.transforms.CenterCrop(224),
            torchvision.transforms.Normalize(
                mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711)
            )
        ])

        if self.config.training:
            transforms = torchvision.transforms.Compose([
                torchvision.transforms.ToTensor(),
                torchvision.transforms.Resize(self.config.image_size[-1], interpolation=torchvision.transforms.InterpolationMode.BILINEAR, antialias=True),
                SmartCrop(self.config.image_size, randomize_p=0.3, randomize_q=0.2)
            ])
        else:
            transforms = None
        controlnet_filter = getattr(controlnet_filters, self.config.controlnet_filter)(
            self.device,
            **(self.config.controlnet_filter_params if self.config.controlnet_filter_params is not None else {})
        )

        return self.Extras(
            gdf=gdf,
            sampling_configs=sampling_configs,
            transforms=transforms,
            effnet_preprocess=effnet_preprocess,
            clip_preprocess=clip_preprocess,
            controlnet_filter=controlnet_filter
        )
    def get_cnet(self, batch: dict, models: Models, extras: Extras, cnet_input=None, target_size=None, **kwargs):
        images = batch['images']
        if target_size is not None:
            images = Image.resize(images, target_size)
        with torch.no_grad():
            if cnet_input is None:
                cnet_input = extras.controlnet_filter(images, **kwargs)
            if isinstance(cnet_input, tuple):
                cnet_input, cnet_input_preview = cnet_input
            else:
                cnet_input_preview = cnet_input
            cnet_input, cnet_input_preview = cnet_input.to(self.device), cnet_input_preview.to(self.device)
        cnet = models.controlnet(cnet_input)
        return cnet, cnet_input_preview

    def get_conditions(self, batch: dict, models: Models, extras: Extras, is_eval=False, is_unconditional=False,

                       eval_image_embeds=False, return_fields=None):
        conditions = super().get_conditions(
            batch, models, extras, is_eval, is_unconditional,
            eval_image_embeds, return_fields=return_fields or ['clip_text', 'clip_text_pooled', 'clip_img']
        )
        return conditions

    def setup_models(self, extras: Extras) -> Models:   # configure model

        
        dtype = getattr(torch, self.config.dtype) if self.config.dtype else torch.bfloat16

        # EfficientNet encoderin
        effnet = EfficientNetEncoder()
        effnet_checkpoint = load_or_fail(self.config.effnet_checkpoint_path)
        effnet.load_state_dict(effnet_checkpoint if 'state_dict' not in effnet_checkpoint else effnet_checkpoint['state_dict'])
        effnet.eval().requires_grad_(False).to(self.device)
        del effnet_checkpoint

        # Previewer
        previewer = Previewer()
        previewer_checkpoint = load_or_fail(self.config.previewer_checkpoint_path)
        previewer.load_state_dict(previewer_checkpoint if 'state_dict' not in previewer_checkpoint else previewer_checkpoint['state_dict'])
        previewer.eval().requires_grad_(False).to(self.device)
        del previewer_checkpoint

        @contextmanager
        def dummy_context():
            yield None

        loading_context = dummy_context if self.config.training else init_empty_weights

        # Diffusion models
        with loading_context():
            generator_ema = None
            if self.config.model_version == '3.6B':
                generator = StageC()
                if self.config.ema_start_iters is not None:  # default setting
                    generator_ema = StageC()
            elif self.config.model_version == '1B':
               
                generator = StageC(c_cond=1536, c_hidden=[1536, 1536], nhead=[24, 24], blocks=[[4, 12], [12, 4]])
                
                if self.config.ema_start_iters is not None and self.config.training:
                    generator_ema = StageC(c_cond=1536, c_hidden=[1536, 1536], nhead=[24, 24], blocks=[[4, 12], [12, 4]])
            else:
                raise ValueError(f"Unknown model version {self.config.model_version}")

       
        
        if loading_context is dummy_context:
            generator.load_state_dict( load_or_fail(self.config.generator_checkpoint_path))
        else:
            for param_name, param in load_or_fail(self.config.generator_checkpoint_path).items():
                    set_module_tensor_to_device(generator, param_name, "cpu", value=param)

        generator._init_extra_parameter()
        

        
       
        generator = generator.to(torch.bfloat16).to(self.device)
     
        train_norm = nn.ModuleList()
        
        
        cnt_norm = 0
        for mm in generator.modules():
            if isinstance(mm,  GlobalResponseNorm):
               
                train_norm.append(Null_Model())
                cnt_norm += 1
                
        
        
        
        train_norm.append(generator.agg_net)
        train_norm.append(generator.agg_net_up)
        
       

        
        if os.path.exists(os.path.join(self.config.output_path, self.config.experiment_id, 'train_norm.safetensors')):
            sdd = torch.load(os.path.join(self.config.output_path, self.config.experiment_id, 'train_norm.safetensors'), map_location='cpu')
            collect_sd = {}
            for k, v in sdd.items():
                collect_sd[k[7:]] = v
            train_norm.load_state_dict(collect_sd, strict=True)
        
        
        train_norm.to(self.device).train().requires_grad_(True)
        train_norm_ema = copy.deepcopy(train_norm)
        train_norm_ema.to(self.device).eval().requires_grad_(False)
        if generator_ema is not None:
           
            generator_ema.load_state_dict(load_or_fail(self.config.generator_checkpoint_path))
            generator_ema._init_extra_parameter()
           
            pretrained_pth = os.path.join(self.config.output_path, self.config.experiment_id, 'generator.safetensors')
            if os.path.exists(pretrained_pth):
              print(pretrained_pth, 'exists')
              generator_ema.load_state_dict(torch.load(pretrained_pth, map_location='cpu'))
          
            generator_ema.eval().requires_grad_(False)
          
        check_nan_inmodel(generator, 'generator')
     
        
        
        if self.config.use_fsdp and self.config.training:
            train_norm = DDP(train_norm, device_ids=[self.device], find_unused_parameters=True)

       
        # CLIP encoders
        tokenizer = AutoTokenizer.from_pretrained(self.config.clip_text_model_name)
        text_model = CLIPTextModelWithProjection.from_pretrained(self.config.clip_text_model_name).requires_grad_(False).to(dtype).to(self.device)
        image_model = CLIPVisionModelWithProjection.from_pretrained(self.config.clip_image_model_name).requires_grad_(False).to(dtype).to(self.device)

        controlnet = ControlNet(
            c_in=extras.controlnet_filter.num_channels(),
            proj_blocks=self.config.controlnet_blocks,
            bottleneck_mode=self.config.controlnet_bottleneck_mode
        )
        controlnet = controlnet.to(dtype).to(self.device)
        controlnet = self.load_model(controlnet, 'controlnet')
        controlnet.backbone.eval().requires_grad_(True)
        
        
        return self.Models(
            effnet=effnet, previewer=previewer, train_norm = train_norm,
            generator=generator, generator_ema=generator_ema,
            tokenizer=tokenizer, text_model=text_model, image_model=image_model,
            train_norm_ema=train_norm_ema, controlnet =controlnet
        )

    def setup_optimizers(self, extras: Extras, models: Models) -> TrainingCore.Optimizers:
        
#
      
        params = []
        params += list(models.train_norm.module.parameters())

        optimizer = optim.AdamW(params, lr=self.config.lr) 

        return self.Optimizers(generator=optimizer)

    def ema_update(self, ema_model, source_model, beta):
        for param_src, param_ema in zip(source_model.parameters(), ema_model.parameters()):
            param_ema.data.mul_(beta).add_(param_src.data, alpha = 1 - beta)
            
    def sync_ema(self, ema_model):
        print('sync ema', torch.distributed.get_world_size())
        for param in ema_model.parameters():
            torch.distributed.all_reduce(param.data, op=torch.distributed.ReduceOp.SUM)
            param.data /= torch.distributed.get_world_size()
    def setup_optimizers_backup(self, extras: Extras, models: Models) -> TrainingCore.Optimizers:
       

        optimizer = optim.AdamW(
            models.generator.up_blocks.parameters() , 
        lr=self.config.lr)
        optimizer = self.load_optimizer(optimizer, 'generator_optim',
                                        fsdp_model=models.generator if self.config.use_fsdp else None)
        return self.Optimizers(generator=optimizer)

    def setup_schedulers(self, extras: Extras, models: Models, optimizers: TrainingCore.Optimizers) -> Schedulers:
        scheduler = GradualWarmupScheduler(optimizers.generator, multiplier=1, total_epoch=self.config.warmup_updates)
        scheduler.last_epoch = self.info.total_steps
        return self.Schedulers(generator=scheduler)

    def setup_data(self, extras: Extras) -> WarpCore.Data:
        # SETUP DATASET
        dataset_path = self.config.webdataset_path
        print('in line 96', dataset_path, type(dataset_path))

        dataset = mydist_dataset(dataset_path, \
            torchvision.transforms.ToTensor() if self.config.multi_aspect_ratio is not None \
                else extras.transforms)

        # SETUP DATALOADER
        real_batch_size = self.config.batch_size // (self.world_size * self.config.grad_accum_steps)
        print('in line 119', self.process_id, real_batch_size)
        sampler =  DistributedSampler(dataset, rank=self.process_id, num_replicas = self.world_size, shuffle=True)
        dataloader = DataLoader(
            dataset, batch_size=real_batch_size, num_workers=4, pin_memory=True,
            collate_fn=identity if self.config.multi_aspect_ratio is not None else None,
            sampler = sampler
        )
        if self.is_main_node:
            print(f"Training with batch size {self.config.batch_size} ({real_batch_size}/GPU)")

        if self.config.multi_aspect_ratio is not None:
            aspect_ratios = [float(Fraction(f)) for f in self.config.multi_aspect_ratio]
            dataloader_iterator = Bucketeer(dataloader, density=[ss*ss for ss in self.config.image_size] , factor=32,
                                            ratios=aspect_ratios, p_random_ratio=self.config.bucketeer_random_ratio,
                                            interpolate_nearest=False)  # , use_smartcrop=True)
        else:
           
            dataloader_iterator = iter(dataloader)

        return self.Data(dataset=dataset, dataloader=dataloader, iterator=dataloader_iterator, sampler=sampler)





    def setup_ddp(self, experiment_id, single_gpu=False, rank=0):

        if not single_gpu:
            local_rank = rank
            process_id = rank
            world_size = get_world_size()

            self.process_id = process_id
            self.is_main_node = process_id == 0
            self.device = torch.device(local_rank)
            self.world_size = world_size

          
            os.environ['MASTER_ADDR'] = 'localhost'
            os.environ['MASTER_PORT'] = '41443'
            torch.cuda.set_device(local_rank)
            init_process_group(
                backend="nccl",
                rank=local_rank,
                world_size=world_size,
                # init_method=init_method,
            )
            print(f"[GPU {process_id}] READY")
        else:
            self.is_main_node = rank == 0
            self.process_id = rank
            self.device = torch.device('cuda:0')
            self.world_size = 1
            print("Running in single thread, DDP not enabled.")
    # Training loop --------------------------------
    def get_target_lr_size(self, ratio, std_size=24):
        w, h = int(std_size / math.sqrt(ratio)), int(std_size * math.sqrt(ratio)) 
        return (h * 32 , w * 32) 
    def forward_pass(self, data: WarpCore.Data, extras: Extras, models: Models):
        #batch = next(data.iterator)
        batch = data
        ratio = batch['images'].shape[-2] / batch['images'].shape[-1]
        shape_lr = self.get_target_lr_size(ratio)

        with torch.no_grad():
            conditions = self.get_conditions(batch, models, extras)
            
            latents = self.encode_latents(batch, models, extras)
            latents_lr = self.encode_latents(batch, models, extras,target_size=shape_lr)
           
            noised, noise, target, logSNR, noise_cond, loss_weight = extras.gdf.diffuse(latents, shift=1, loss_shift=1)
            noised_lr, noise_lr, target_lr, logSNR_lr, noise_cond_lr, loss_weight_lr = extras.gdf.diffuse(latents_lr, shift=1, loss_shift=1, t=torch.ones(latents.shape[0]).to(latents.device)*0.05, )

        with torch.cuda.amp.autocast(dtype=torch.bfloat16):
            
            require_cond = True
          
            with torch.no_grad():
                _, lr_enc_guide, lr_dec_guide = models.generator(noised_lr, noise_cond_lr, reuire_f=True, **conditions)
            
            
            pred = models.generator(noised, noise_cond, reuire_f=False, lr_guide=(lr_enc_guide, lr_dec_guide) if require_cond else None , **conditions)             
            loss = nn.functional.mse_loss(pred, target, reduction='none').mean(dim=[1, 2, 3]) 
           
            loss_adjusted = (loss * loss_weight ).mean() / self.config.grad_accum_steps 
            #
        if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
            extras.gdf.loss_weight.update_buckets(logSNR, loss)
       
        return loss,  loss_adjusted

    def backward_pass(self, update, loss_adjusted, models: Models, optimizers: TrainingCore.Optimizers, schedulers: Schedulers):
     
        if update:
        
            torch.distributed.barrier()
            loss_adjusted.backward()
            
            
            grad_norm = nn.utils.clip_grad_norm_(models.train_norm.module.parameters(), 1.0)
          
            optimizers_dict = optimizers.to_dict()
            for k in optimizers_dict:
                if k != 'training':
                    optimizers_dict[k].step()
            schedulers_dict = schedulers.to_dict()
            for k in schedulers_dict:
                if k != 'training':
                    schedulers_dict[k].step()
            for k in optimizers_dict:
                if k != 'training':
                    optimizers_dict[k].zero_grad(set_to_none=True)
            self.info.total_steps += 1
        else:
            #print('in line 457', loss_adjusted)
            loss_adjusted.backward()
            #torch.distributed.barrier()
            grad_norm = torch.tensor(0.0).to(self.device)
        
        return grad_norm

    def models_to_save(self):
        return ['generator', 'generator_ema', 'trans_inr', 'trans_inr_ema']

    def encode_latents(self, batch: dict, models: Models, extras: Extras, target_size=None) -> torch.Tensor:
        
        images = batch['images'].to(self.device)
        if target_size is not None:
          images = F.interpolate(images, target_size)
          #images = apply_degradations(images)
        return models.effnet(extras.effnet_preprocess(images))

    def decode_latents(self, latents: torch.Tensor, batch: dict, models: Models, extras: Extras) -> torch.Tensor:
        return models.previewer(latents)

    def __init__(self, rank=0, config_file_path=None, config_dict=None, device="cpu", training=True, world_size=1, ):
        # Temporary setup, will be overriden by setup_ddp if required
        # self.device = device
        # self.process_id = 0
        # self.is_main_node = True
        # self.world_size = 1
        # ----
        # self.world_size = world_size
        # self.process_id = rank
        # self.device=device
        self.is_main_node = (rank == 0)
        self.config: self.Config = self.setup_config(config_file_path, config_dict, training)
        self.setup_ddp(self.config.experiment_id, single_gpu=world_size <= 1, rank=rank)
        self.info: self.Info = self.setup_info()
        print('in line 292', self.config.experiment_id, rank, world_size <= 1)
        p = [i for i in range( 2 * 768 // 32)]
        p = [num / sum(p) for num in p]
        self.rand_pro = p
        self.res_list = [o for o in range(800, 2336, 32)]
        
        #[32, 40, 48]
        #in line 292 stage_c_3b_finetuning False
        
    def __call__(self, single_gpu=False):
         # this will change the device to the CUDA rank
        #self.setup_wandb()
        if self.config.allow_tf32:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True

        if self.is_main_node:
            print()
            print("**STARTIG JOB WITH CONFIG:**")
            print(yaml.dump(self.config.to_dict(), default_flow_style=False))
            print("------------------------------------")
            print()
            print("**INFO:**")
            print(yaml.dump(vars(self.info), default_flow_style=False))
            print("------------------------------------")
            print()
        print('in line 308', self.is_main_node, self.is_main_node, self.process_id, self.device  )
        # SETUP STUFF
        extras = self.setup_extras_pre()
        assert extras is not None, "setup_extras_pre() must return a DTO"



        data = self.setup_data(extras)
        assert data is not None, "setup_data() must return a DTO"
        if self.is_main_node:
            print("**DATA:**")
            print(yaml.dump({k:type(v).__name__ for k, v in data.to_dict().items()}, default_flow_style=False))
            print("------------------------------------")
            print()

        models = self.setup_models(extras)
        assert models is not None, "setup_models() must return a DTO"
        if self.is_main_node:
            print("**MODELS:**")
            print(yaml.dump({
                k:f"{type(v).__name__} - {f'trainable params {sum(p.numel() for p in v.parameters() if p.requires_grad)}' if isinstance(v, nn.Module) else 'Not a nn.Module'}" for k, v in models.to_dict().items()
            }, default_flow_style=False))
            print("------------------------------------")
            print()



        optimizers = self.setup_optimizers(extras, models)
        assert optimizers is not None, "setup_optimizers() must return a DTO"
        if self.is_main_node:
            print("**OPTIMIZERS:**")
            print(yaml.dump({k:type(v).__name__ for k, v in optimizers.to_dict().items()}, default_flow_style=False))
            print("------------------------------------")
            print()

        schedulers = self.setup_schedulers(extras, models, optimizers)
        assert schedulers is not None, "setup_schedulers() must return a DTO"
        if self.is_main_node:
            print("**SCHEDULERS:**")
            print(yaml.dump({k:type(v).__name__ for k, v in schedulers.to_dict().items()}, default_flow_style=False))
            print("------------------------------------")
            print()

        post_extras =self.setup_extras_post(extras, models, optimizers, schedulers)
        assert post_extras is not None, "setup_extras_post() must return a DTO"
        extras = self.Extras.from_dict({ **extras.to_dict(),**post_extras.to_dict() })
        if self.is_main_node:
            print("**EXTRAS:**")
            print(yaml.dump({k:f"{v}" for k, v in extras.to_dict().items()}, default_flow_style=False))
            print("------------------------------------")
            print()
        # -------

        # TRAIN
        if self.is_main_node:
            print("**TRAINING STARTING...**")
        self.train(data, extras, models, optimizers, schedulers)

        if single_gpu is False:
            barrier()
            destroy_process_group()
        if self.is_main_node:
            print()
            print("------------------------------------")
            print()
            print("**TRAINING COMPLETE**")
            if self.config.wandb_project is not None:
                wandb.alert(title=f"Training {self.info.wandb_run_id} finished", text=f"Training {self.info.wandb_run_id} finished")


    def train(self, data: WarpCore.Data, extras: WarpCore.Extras, models: Models, optimizers: TrainingCore.Optimizers,

              schedulers: WarpCore.Schedulers):
        start_iter = self.info.iter + 1
        max_iters = self.config.updates * self.config.grad_accum_steps
        if self.is_main_node:
            print(f"STARTING AT STEP: {start_iter}/{max_iters}")

     
        if self.is_main_node:
            create_folder_if_necessary(f'{self.config.output_path}/{self.config.experiment_id}/')
        if 'generator' in self.models_to_save():
            models.generator.train()
        #initial_params = {name: param.clone() for name, param in models.train_norm.named_parameters()}
        iter_cnt = 0
        epoch_cnt = 0
        models.train_norm.train()
        while True:
          epoch_cnt += 1
          if self.world_size > 1:
            print('sampler set epoch', epoch_cnt)
            data.sampler.set_epoch(epoch_cnt)  
          for ggg in range(len(data.dataloader)):
              iter_cnt += 1
              # FORWARD PASS
              #print('in line 414 before forward', iter_cnt, batch['captions'][0], self.process_id)
              #loss, loss_adjusted, loss_extra = self.forward_pass(batch, extras, models)
              loss, loss_adjusted = self.forward_pass(next(data.iterator), extras, models)
              
              #print('in line 416', loss, iter_cnt)
              # # BACKWARD PASS
    
              grad_norm = self.backward_pass(
                        iter_cnt % self.config.grad_accum_steps == 0 or iter_cnt == max_iters, loss_adjusted,
                        models, optimizers, schedulers
                      )
              
              
              
              self.info.iter = iter_cnt
              
              # UPDATE EMA
              if  iter_cnt % self.config.ema_iters == 0:
                 
                  with torch.no_grad():
                      print('in line 890 ema update', self.config.ema_iters, iter_cnt)
                      self.ema_update(models.train_norm_ema, models.train_norm, self.config.ema_beta)
                      #generator.module.agg_net.
                      #self.ema_update(models.generator_ema.agg_net, models.generator.module.agg_net, self.config.ema_beta)
                      #self.ema_update(models.generator_ema.agg_net_up, models.generator.module.agg_net_up, self.config.ema_beta)
                      
              # UPDATE LOSS METRICS
              self.info.ema_loss = loss.mean().item() if self.info.ema_loss is None else self.info.ema_loss * 0.99 + loss.mean().item() * 0.01
  
              #print('in line 666 after ema loss', grad_norm, loss.mean().item(), iter_cnt, self.info.ema_loss)
              if self.is_main_node and  np.isnan(loss.mean().item()) or np.isnan(grad_norm.item()):
                      print(f"gggg NaN value encountered in training run {self.info.wandb_run_id}", \
                      f"Loss {loss.mean().item()} - Grad Norm {grad_norm.item()}. Run {self.info.wandb_run_id}")
  
              if self.is_main_node:
                  logs = {
                      'loss': self.info.ema_loss,
                      'backward_loss': loss_adjusted.mean().item(),
                      #'raw_extra_loss': loss_extra.mean().item(),
                      'ema_loss': self.info.ema_loss,
                      'raw_ori_loss': loss.mean().item(),
                      #'raw_rec_loss': loss_rec.mean().item(),
                      #'raw_lr_loss': loss_lr.mean().item(),
                      #'reg_loss':loss_reg.item(),
                      'grad_norm': grad_norm.item(),
                      'lr': optimizers.generator.param_groups[0]['lr'] if optimizers.generator is not None else 0,
                      'total_steps': self.info.total_steps,
                  }
                  if iter_cnt % (self.config.save_every) == 0:
                        
                      print(iter_cnt, max_iters, logs, epoch_cnt, )
                  #pbar.set_postfix(logs)
                 
  
              #if iter_cnt % 10 == 0:
                  
              
              if iter_cnt == 1 or iter_cnt % (self.config.save_every  ) == 0 or iter_cnt == max_iters:
              #if True:
                  # SAVE AND CHECKPOINT STUFF
                  if np.isnan(loss.mean().item()):
                      if self.is_main_node and self.config.wandb_project is not None:
                          print(f"NaN value encountered in training run {self.info.wandb_run_id}", \
                          f"Loss {loss.mean().item()} - Grad Norm {grad_norm.item()}. Run {self.info.wandb_run_id}")
                     
                  else:
                      if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
                          self.info.adaptive_loss = {
                              'bucket_ranges': extras.gdf.loss_weight.bucket_ranges.tolist(),
                              'bucket_losses': extras.gdf.loss_weight.bucket_losses.tolist(),
                          }
                      #self.save_checkpoints(models, optimizers)
           
                      #torch.save(models.trans_inr.module.state_dict(), \
                      #f'{self.config.output_path}/{self.config.experiment_id}/trans_inr.safetensors')
                      #torch.save(models.trans_inr_ema.state_dict(), \
                      #f'{self.config.output_path}/{self.config.experiment_id}/trans_inr_ema.safetensors')
                      
                      
                      if self.is_main_node and iter_cnt % (self.config.save_every * self.config.grad_accum_steps) == 0:
                          print('save model', iter_cnt, iter_cnt % (self.config.save_every * self.config.grad_accum_steps), self.config.save_every, self.config.grad_accum_steps )
                          torch.save(models.train_norm.state_dict(), \
                          f'{self.config.output_path}/{self.config.experiment_id}/train_norm.safetensors')
                          
                          #self.sync_ema(models.train_norm_ema)
                          torch.save(models.train_norm_ema.state_dict(), \
                          f'{self.config.output_path}/{self.config.experiment_id}/train_norm_ema.safetensors')
                          #if self.is_main_node and iter_cnt % (4 * self.config.save_every * self.config.grad_accum_steps) == 0:
                          torch.save(models.train_norm.state_dict(), \
                              f'{self.config.output_path}/{self.config.experiment_id}/train_norm_{iter_cnt}.safetensors')
                          
                       
              if iter_cnt == 1 or iter_cnt % (self.config.save_every* self.config.grad_accum_steps) == 0 or iter_cnt == max_iters:
                  
                  if self.is_main_node:
                     #check_nan_inmodel(models.generator, 'generator')
                     #check_nan_inmodel(models.generator_ema,  'generator_ema')
                     self.sample(models, data, extras)
              if False:
                param_changes = {name: (param - initial_params[name]).norm().item() for name, param in models.train_norm.named_parameters()}
                threshold = sorted(param_changes.values(), reverse=True)[int(len(param_changes) * 0.1)]  # top 10%
                important_params = [name for name, change in param_changes.items() if change > threshold]
                print(important_params, threshold, len(param_changes), self.process_id)
                json.dump(important_params, open(f'{self.config.output_path}/{self.config.experiment_id}/param.json', 'w'), indent=4)     
                  
         
          if self.info.iter >= max_iters:
            break
            
    def sample(self, models: Models, data: WarpCore.Data, extras: Extras):
       
        #if 'generator' in self.models_to_save():
        models.generator.eval()
        models.train_norm.eval()
        with torch.no_grad():
            batch = next(data.iterator)
            ratio = batch['images'].shape[-2] / batch['images'].shape[-1]
            #batch['images'] = batch['images'].to(torch.float16)
            shape_lr = self.get_target_lr_size(ratio)
            conditions = self.get_conditions(batch, models, extras, is_eval=True, is_unconditional=False, eval_image_embeds=False)
            unconditions = self.get_conditions(batch, models, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)
            cnet, cnet_input = self.get_cnet(batch, models, extras)
            conditions, unconditions = {**conditions, 'cnet': cnet}, {**unconditions, 'cnet': cnet}
            
            latents = self.encode_latents(batch, models, extras)
            latents_lr = self.encode_latents(batch, models, extras, target_size = shape_lr)
           
            if self.is_main_node:
                
                with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                    #print('in line 366 on v100 switch to tf16')
                    *_, (sampled, _, _, sampled_lr) = extras.gdf.sample(
                        models.generator, models.trans_inr, conditions,
                        latents.shape, latents_lr.shape, 
                        unconditions, device=self.device, **extras.sampling_configs
                    )
    
                  
                       
                    #else:
                    sampled_ema = sampled
                    sampled_ema_lr = sampled_lr
            
            
            if self.is_main_node:
                print('sampling results', latents.shape, latents_lr.shape, )
                noised_images = torch.cat(
                    [self.decode_latents(latents[i:i + 1].float(), batch, models, extras) for i in range(len(latents))], dim=0)
                
                sampled_images = torch.cat(
                    [self.decode_latents(sampled[i:i + 1].float(), batch, models, extras) for i in range(len(sampled))], dim=0)
                sampled_images_ema = torch.cat(
                    [self.decode_latents(sampled_ema[i:i + 1].float(), batch, models, extras) for i in range(len(sampled_ema))],
                    dim=0)
                    
                noised_images_lr = torch.cat(
                    [self.decode_latents(latents_lr[i:i + 1].float(), batch, models, extras) for i in range(len(latents_lr))], dim=0)
                
                sampled_images_lr = torch.cat(
                    [self.decode_latents(sampled_lr[i:i + 1].float(), batch, models, extras) for i in range(len(sampled_lr))], dim=0)
                sampled_images_ema_lr = torch.cat(
                    [self.decode_latents(sampled_ema_lr[i:i + 1].float(), batch, models, extras) for i in range(len(sampled_ema_lr))],
                    dim=0)

                images = batch['images']
                if images.size(-1) != noised_images.size(-1) or images.size(-2) != noised_images.size(-2):
                    images = nn.functional.interpolate(images, size=noised_images.shape[-2:], mode='bicubic')
                    images_lr = nn.functional.interpolate(images, size=noised_images_lr.shape[-2:], mode='bicubic')

                collage_img = torch.cat([
                    torch.cat([i for i in images.cpu()], dim=-1),
                    torch.cat([i for i in noised_images.cpu()], dim=-1),
                    torch.cat([i for i in sampled_images.cpu()], dim=-1),
                    torch.cat([i for i in sampled_images_ema.cpu()], dim=-1),
                ], dim=-2)
                
                collage_img_lr = torch.cat([
                    torch.cat([i for i in images_lr.cpu()], dim=-1),
                    torch.cat([i for i in noised_images_lr.cpu()], dim=-1),
                    torch.cat([i for i in sampled_images_lr.cpu()], dim=-1),
                    torch.cat([i for i in sampled_images_ema_lr.cpu()], dim=-1),
                ], dim=-2)

                torchvision.utils.save_image(collage_img, f'{self.config.output_path}/{self.config.experiment_id}/{self.info.total_steps:06d}.jpg')
                torchvision.utils.save_image(collage_img_lr, f'{self.config.output_path}/{self.config.experiment_id}/{self.info.total_steps:06d}_lr.jpg')
                #torchvision.utils.save_image(collage_img, f'{self.config.experiment_id}_latest_output.jpg')

                captions = batch['captions']
                if self.config.wandb_project is not None:
                    log_data = [
                        [captions[i]] + [wandb.Image(sampled_images[i])] + [wandb.Image(sampled_images_ema[i])] + [
                            wandb.Image(images[i])] for i in range(len(images))]
                    log_table = wandb.Table(data=log_data, columns=["Captions", "Sampled", "Sampled EMA", "Orig"])
                    wandb.log({"Log": log_table})

                    if isinstance(extras.gdf.loss_weight, AdaptiveLossWeight):
                        plt.plot(extras.gdf.loss_weight.bucket_ranges, extras.gdf.loss_weight.bucket_losses[:-1])
                        plt.ylabel('Raw Loss')
                        plt.ylabel('LogSNR')
                        wandb.log({"Loss/LogSRN": plt})

            #if 'generator' in self.models_to_save():
            models.generator.train()
            models.train_norm.train()
            print('finishe sampling in line 901')
    
    
    
    def sample_fortest(self, models: Models, extras: Extras, hr_shape, lr_shape, batch, eval_image_embeds=False):
       
        #if 'generator' in self.models_to_save():
        models.generator.eval()
        models.trans_inr.eval()
        models.controlnet.eval()
        with torch.no_grad():
           
            if self.is_main_node:
                conditions = self.get_conditions(batch, models, extras, is_eval=True, is_unconditional=False, eval_image_embeds=eval_image_embeds)
                unconditions = self.get_conditions(batch, models, extras, is_eval=True, is_unconditional=True, eval_image_embeds=False)
                cnet, cnet_input = self.get_cnet(batch, models, extras, target_size = lr_shape)
                conditions, unconditions = {**conditions, 'cnet': cnet}, {**unconditions, 'cnet': cnet}
                
                #print('in line 885', self.is_main_node)
                with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                    #print('in line 366 on v100 switch to tf16')
                    *_, (sampled, _, _, sampled_lr) = extras.gdf.sample(
                        models.generator, models.trans_inr, conditions,
                        hr_shape, lr_shape, 
                        unconditions, device=self.device, **extras.sampling_configs
                    )
    
                    if models.generator_ema is not None:
                        
                        *_, (sampled_ema, _, _, sampled_ema_lr) = extras.gdf.sample(
                            models.generator_ema,  models.trans_inr_ema,  conditions,
                            latents.shape, latents_lr.shape, 
                            unconditions, device=self.device, **extras.sampling_configs
                        )
                       
                    else:
                        sampled_ema = sampled
                        sampled_ema_lr = sampled_lr
            #x0, x, epsilon, x0_lr, x_lr, pred_lr)
            #sampled, _ = models.trans_inr(sampled, None, sampled)
            #sampled_lr, _ = models.trans_inr(sampled, None, sampled_lr)
            
        return sampled, sampled_lr
def main_worker(rank, cfg):
    print("Launching Script in main worker")
    print('in line 467', rank)
    warpcore = WurstCore(
        config_file_path=cfg, rank=rank, world_size = get_world_size()
    )
    # core.fsdp_defaults['sharding_strategy'] = ShardingStrategy.NO_SHARD

    # RUN TRAINING
    warpcore(get_world_size()==1)

if __name__ == '__main__':
    print('launch multi process')
    # os.environ["OMP_NUM_THREADS"] = "1" 
    # os.environ["MKL_NUM_THREADS"] = "1" 
    #dist.init_process_group(backend="nccl")
    #torch.backends.cudnn.benchmark = True
#train/train_c_my.py
    #mp.set_sharing_strategy('file_system')
    print('in line 481', sys.argv[1] if len(sys.argv) > 1 else None)
    print('in line 481',get_master_ip(), get_world_size() )
    print('in line 484', get_world_size())
    if get_master_ip() == "127.0.0.1":
        # manually launch distributed processes
        mp.spawn(main_worker, nprocs=get_world_size(), args=(sys.argv[1] if len(sys.argv) > 1 else None, ))
    else:
        main_worker(0, sys.argv[1] if len(sys.argv) > 1 else None, )