# Copyright 2024 The HuggingFace Team. # All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import signal import sys import threading import time import cv2 sys.path.append('..') from PIL import Image import logging import math import os from pathlib import Path import torch import transformers from accelerate import Accelerator from accelerate.logging import get_logger from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed from huggingface_hub import create_repo from torch.utils.data import DataLoader from tqdm.auto import tqdm import numpy as np from transformers import AutoTokenizer, T5EncoderModel import diffusers from diffusers import AutoencoderKLCogVideoX, CogVideoXDPMScheduler from diffusers.optimization import get_scheduler from diffusers.training_utils import ( cast_training_params, free_memory, ) from diffusers.utils import check_min_version, export_to_video, is_wandb_available from diffusers.utils.torch_utils import is_compiled_module from controlnet_datasets import FullMotionBlurDataset, GoPro2xMotionBlurDataset, OutsidePhotosDataset, GoProMotionBlurDataset, BAISTDataset from controlnet_pipeline import ControlnetCogVideoXPipeline from cogvideo_transformer import CogVideoXTransformer3DModel from helpers import random_insert_latent_frame, transform_intervals import os from utils import save_frames_as_pngs, compute_prompt_embeddings, prepare_rotary_positional_embeddings, encode_prompt, get_optimizer, atomic_save, get_args if is_wandb_available(): import wandb # Will error if the minimal version of diffusers is not installed. Remove at your own risks. check_min_version("0.31.0.dev0") logger = get_logger(__name__) def log_validation( pipe, args, accelerator, pipeline_args, ): logger.info( f"Running validation... \n Generating {args.num_validation_videos} videos with prompt: {pipeline_args['prompt']}." ) # We train on the simplified learning objective. If we were previously predicting a variance, we need the scheduler to ignore it scheduler_args = {} if "variance_type" in pipe.scheduler.config: variance_type = pipe.scheduler.config.variance_type if variance_type in ["learned", "learned_range"]: variance_type = "fixed_small" scheduler_args["variance_type"] = variance_type pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, **scheduler_args) pipe = pipe.to(accelerator.device) # run inference generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None videos = [] for _ in range(args.num_validation_videos): video = pipe(**pipeline_args, generator=generator, output_type="np").frames[0] videos.append(video) free_memory() #delete the pipeline to free up memory return videos def main(args): global signal_recieved_time if args.report_to == "wandb" and args.hub_token is not None: raise ValueError( "You cannot use both --report_to=wandb and --hub_token due to a security risk of exposing your token." " Please use `huggingface-cli login` to authenticate with the Hub." ) if torch.backends.mps.is_available() and args.mixed_precision == "bf16": # due to pytorch#99272, MPS does not yet support bfloat16. raise ValueError( "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." ) logging_dir = Path(args.output_dir, args.logging_dir) accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, logging_dir=logging_dir) kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) accelerator = Accelerator( gradient_accumulation_steps=args.gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with=args.report_to, project_config=accelerator_project_config, kwargs_handlers=[kwargs], ) # Disable AMP for MPS. if torch.backends.mps.is_available(): accelerator.native_amp = False if args.report_to == "wandb": if not is_wandb_available(): raise ImportError("Make sure to install wandb if you want to use it for logging during training.") # Make one log on every process with the configuration for debugging. logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO, ) logger.info(accelerator.state, main_process_only=False) if accelerator.is_local_main_process: transformers.utils.logging.set_verbosity_warning() diffusers.utils.logging.set_verbosity_info() else: transformers.utils.logging.set_verbosity_error() diffusers.utils.logging.set_verbosity_error() # If passed along, set the training seed now. if args.seed is not None: set_seed(args.seed) # Handle the repository creation if accelerator.is_main_process: if args.output_dir is not None: os.makedirs(args.output_dir, exist_ok=True) if args.push_to_hub: repo_id = create_repo( repo_id=args.hub_model_id or Path(args.output_dir).name, exist_ok=True, ).repo_id # Prepare models and scheduler tokenizer = AutoTokenizer.from_pretrained( os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="tokenizer", revision=args.revision ) text_encoder = T5EncoderModel.from_pretrained( os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="text_encoder", revision=args.revision ) # CogVideoX-2b weights are stored in float16 config = CogVideoXTransformer3DModel.load_config( os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="transformer", revision=args.revision, variant=args.variant, ) load_dtype = torch.bfloat16 if "5b" in os.path.join(args.base_dir, args.pretrained_model_name_or_path).lower() else torch.float16 transformer = CogVideoXTransformer3DModel.from_pretrained( os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="transformer", torch_dtype=load_dtype, revision=args.revision, variant=args.variant, low_cpu_mem_usage=False, ) vae = AutoencoderKLCogVideoX.from_pretrained( os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="vae", revision=args.revision, variant=args.variant ) scheduler = CogVideoXDPMScheduler.from_pretrained(os.path.join(args.base_dir, args.pretrained_model_name_or_path), subfolder="scheduler") if args.enable_slicing: vae.enable_slicing() if args.enable_tiling: vae.enable_tiling() # We only train the additional adapter controlnet layers text_encoder.requires_grad_(False) transformer.requires_grad_(True) vae.requires_grad_(False) # For mixed precision training we cast all non-trainable weights (vae, text_encoder and transformer) to half-precision # as these weights are only used for inference, keeping weights in full precision is not required. weight_dtype = torch.float32 if accelerator.state.deepspeed_plugin: # DeepSpeed is handling precision, use what's in the DeepSpeed config if ( "fp16" in accelerator.state.deepspeed_plugin.deepspeed_config and accelerator.state.deepspeed_plugin.deepspeed_config["fp16"]["enabled"] ): weight_dtype = torch.float16 if ( "bf16" in accelerator.state.deepspeed_plugin.deepspeed_config and accelerator.state.deepspeed_plugin.deepspeed_config["bf16"]["enabled"] ): weight_dtype = torch.float16 else: if accelerator.mixed_precision == "fp16": weight_dtype = torch.float16 elif accelerator.mixed_precision == "bf16": weight_dtype = torch.bfloat16 if torch.backends.mps.is_available() and weight_dtype == torch.bfloat16: # due to pytorch#99272, MPS does not yet support bfloat16. raise ValueError( "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead." ) text_encoder.to(accelerator.device, dtype=weight_dtype) transformer.to(accelerator.device, dtype=weight_dtype) vae.to(accelerator.device, dtype=weight_dtype) if args.gradient_checkpointing: transformer.enable_gradient_checkpointing() def unwrap_model(model): model = accelerator.unwrap_model(model) model = model._orig_mod if is_compiled_module(model) else model return model # Enable TF32 for faster training on Ampere GPUs, # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices if args.allow_tf32 and torch.cuda.is_available(): torch.backends.cuda.matmul.allow_tf32 = True if args.scale_lr: args.learning_rate = ( args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes ) # Make sure the trainable params are in float32. if args.mixed_precision == "fp16": # only upcast trainable parameters into fp32 cast_training_params([transformer], dtype=torch.float32) trainable_parameters = list(filter(lambda p: p.requires_grad, transformer.parameters())) # Optimization parameters trainable_parameters_with_lr = {"params": trainable_parameters, "lr": args.learning_rate} params_to_optimize = [trainable_parameters_with_lr] use_deepspeed_optimizer = ( accelerator.state.deepspeed_plugin is not None and "optimizer" in accelerator.state.deepspeed_plugin.deepspeed_config ) use_deepspeed_scheduler = ( accelerator.state.deepspeed_plugin is not None and "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config ) optimizer = get_optimizer(args, params_to_optimize, use_deepspeed=use_deepspeed_optimizer) # Dataset and DataLoader DATASET_REGISTRY = { "gopro": GoProMotionBlurDataset, "gopro2x": GoPro2xMotionBlurDataset, "full": FullMotionBlurDataset, "baist": BAISTDataset, "outsidephotos": OutsidePhotosDataset, # val-only special (no split) } if args.dataset not in DATASET_REGISTRY: raise ValueError(f"Unknown dataset: {args.dataset}") train_dataset_class = DATASET_REGISTRY[args.dataset] val_dataset_class = train_dataset_class common_kwargs = dict( data_dir=os.path.join(args.base_dir, args.video_root_dir), output_dir = args.output_dir, image_size=(args.height, args.width), stride=(args.stride_min, args.stride_max), sample_n_frames=args.max_num_frames, hflip_p=args.hflip_p, ) def build_kwargs(is_train: bool): """Return constructor kwargs, adding split""" kw = dict(common_kwargs) kw["split"] = "train" if is_train else args.val_split return kw train_dataset = train_dataset_class(**build_kwargs(is_train=True)) val_dataset = val_dataset_class(**build_kwargs(is_train=False)) def encode_video(video): video = video.to(accelerator.device, dtype=vae.dtype) video = video.permute(0, 2, 1, 3, 4) # [B, C, F, H, W] latent_dist = vae.encode(video).latent_dist.sample() * vae.config.scaling_factor return latent_dist.permute(0, 2, 1, 3, 4).to(memory_format=torch.contiguous_format) def collate_fn(examples): blur_img = [example["blur_img"] for example in examples] videos = [example["video"] for example in examples] if "high_fps_video" in examples[0]: high_fps_videos = [example["high_fps_video"] for example in examples] high_fps_videos = torch.stack(high_fps_videos) high_fps_videos = high_fps_videos.to(memory_format=torch.contiguous_format).float() if "bbx" in examples[0]: bbx = [example["bbx"] for example in examples] bbx = torch.stack(bbx) bbx = bbx.to(memory_format=torch.contiguous_format).float() prompts = [example["caption"] for example in examples] file_names = [example["file_name"] for example in examples] num_frames = [example["num_frames"] for example in examples] input_intervals = [example["input_interval"] for example in examples] output_intervals = [example["output_interval"] for example in examples] videos = torch.stack(videos) videos = videos.to(memory_format=torch.contiguous_format).float() blur_img = torch.stack(blur_img) blur_img = blur_img.to(memory_format=torch.contiguous_format).float() input_intervals = torch.stack(input_intervals) input_intervals = input_intervals.to(memory_format=torch.contiguous_format).float() output_intervals = torch.stack(output_intervals) output_intervals = output_intervals.to(memory_format=torch.contiguous_format).float() out_dict = { "file_names": file_names, "blur_img": blur_img, "videos": videos, "num_frames": num_frames, "prompts": prompts, "input_intervals": input_intervals, "output_intervals": output_intervals, } if "high_fps_video" in examples[0]: out_dict["high_fps_video"] = high_fps_videos if "bbx" in examples[0]: out_dict["bbx"] = bbx return out_dict train_dataloader = DataLoader( train_dataset, batch_size=args.train_batch_size, shuffle=True, collate_fn=collate_fn, num_workers=args.dataloader_num_workers, ) val_dataloader = DataLoader( val_dataset, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=args.dataloader_num_workers, ) # Scheduler and math around the number of training steps. overrode_max_train_steps = False num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) if args.max_train_steps is None: args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch overrode_max_train_steps = True if use_deepspeed_scheduler: from accelerate.utils import DummyScheduler lr_scheduler = DummyScheduler( name=args.lr_scheduler, optimizer=optimizer, total_num_steps=args.max_train_steps * accelerator.num_processes, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, ) else: lr_scheduler = get_scheduler( args.lr_scheduler, optimizer=optimizer, num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes, num_training_steps=args.max_train_steps * accelerator.num_processes, num_cycles=args.lr_num_cycles, power=args.lr_power, ) # Prepare everything with our `accelerator`. transformer, optimizer, train_dataloader, lr_scheduler, val_dataloader = accelerator.prepare( transformer, optimizer, train_dataloader, lr_scheduler, val_dataloader ) # We need to recalculate our total training steps as the size of the training dataloader may have changed. num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps) if overrode_max_train_steps: args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch # Afterwards we recalculate our number of training epochs args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch) # We need to initialize the trackers we use, and also store our configuration. # The trackers initializes automatically on the main process. if accelerator.is_main_process: tracker_name = args.tracker_name or "cogvideox-controlnet" accelerator.init_trackers(tracker_name, config=vars(args)) accelerator.register_for_checkpointing(transformer, optimizer, lr_scheduler) save_path = os.path.join(args.output_dir, f"checkpoint") #check if the checkpoint already exists if os.path.exists(save_path): accelerator.load_state(save_path) logger.info(f"Loaded state from {save_path}") # Train! total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps num_trainable_parameters = sum(param.numel() for model in params_to_optimize for param in model["params"]) logger.info("***** Running training *****") logger.info(f" Num trainable parameters = {num_trainable_parameters}") logger.info(f" Num examples = {len(train_dataset)}") logger.info(f" Num batches each epoch = {len(train_dataloader)}") logger.info(f" Num epochs = {args.num_train_epochs}") logger.info(f" Instantaneous batch size per device = {args.train_batch_size}") logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}") logger.info(f" Gradient accumulation steps = {args.gradient_accumulation_steps}") logger.info(f" Total optimization steps = {args.max_train_steps}") global_step = 0 first_epoch = 0 initial_global_step = 0 progress_bar = tqdm( range(0, args.max_train_steps), initial=initial_global_step, desc="Steps", # Only show the progress bar once on each machine. disable=not accelerator.is_local_main_process, ) vae_scale_factor_spatial = 2 ** (len(vae.config.block_out_channels) - 1) # For DeepSpeed training model_config = transformer.module.config if hasattr(transformer, "module") else transformer.config for epoch in range(first_epoch, args.num_train_epochs): transformer.train() for step, batch in enumerate(train_dataloader): if not args.just_validate: models_to_accumulate = [transformer] with accelerator.accumulate(models_to_accumulate): model_input = encode_video(batch["videos"]).to(dtype=weight_dtype) # [B, F, C, H, W] prompts = batch["prompts"] image_latent = encode_video(batch["blur_img"]).to(dtype=weight_dtype) # [B, F, C, H, W] input_intervals = batch["input_intervals"] output_intervals = batch["output_intervals"] batch_size = len(prompts) # True = use real prompt (conditional); False = drop to empty (unconditional) guidance_mask = torch.rand(batch_size, device=accelerator.device) >= 0.2 # build a new prompts list: keep the original where mask True, else blank per_sample_prompts = [ prompts[i] if guidance_mask[i] else "" for i in range(batch_size) ] prompts = per_sample_prompts # encode prompts prompt_embeds = compute_prompt_embeddings( tokenizer, text_encoder, prompts, model_config.max_text_seq_length, accelerator.device, weight_dtype, requires_grad=False, ) # Sample noise that will be added to the latents noise = torch.randn_like(model_input) batch_size, num_frames, num_channels, height, width = model_input.shape # Sample a random timestep for each image timesteps = torch.randint( 0, scheduler.config.num_train_timesteps, (batch_size,), device=model_input.device ) timesteps = timesteps.long() # Prepare rotary embeds image_rotary_emb = ( prepare_rotary_positional_embeddings( height=args.height, width=args.width, num_frames=num_frames, vae_scale_factor_spatial=vae_scale_factor_spatial, patch_size=model_config.patch_size, attention_head_dim=model_config.attention_head_dim, device=accelerator.device, ) if model_config.use_rotary_positional_embeddings else None ) # Add noise to the model input according to the noise magnitude at each timestep (this is the forward diffusion process) noisy_model_input = scheduler.add_noise(model_input, noise, timesteps) input_intervals = transform_intervals(input_intervals, frames_per_latent=4) output_intervals = transform_intervals(output_intervals, frames_per_latent=4) #first interval is always rep noisy_model_input, target, condition_mask, intervals = random_insert_latent_frame(image_latent, noisy_model_input, model_input, input_intervals, output_intervals, special_info=args.special_info) for i in range(batch_size): if not guidance_mask[i]: noisy_model_input[i][condition_mask[i]] = 0 # Predict the noise residual model_output = transformer( hidden_states=noisy_model_input, encoder_hidden_states=prompt_embeds, intervals=intervals, condition_mask=condition_mask, timestep=timesteps, image_rotary_emb=image_rotary_emb, return_dict=False, )[0] #this line below is also scaling the input which is bad - so the model is also learning to scale this input latent somehow #thus, we need to replace the first frame with the original frame later model_pred = scheduler.get_velocity(model_output, noisy_model_input, timesteps) alphas_cumprod = scheduler.alphas_cumprod[timesteps] weights = 1 / (1 - alphas_cumprod) while len(weights.shape) < len(model_pred.shape): weights = weights.unsqueeze(-1) loss = torch.mean((weights * (model_pred[~condition_mask] - target[~condition_mask]) ** 2).reshape(batch_size, -1), dim=1) loss = loss.mean() accelerator.backward(loss) if accelerator.state.deepspeed_plugin is None: if not args.just_validate: optimizer.step() optimizer.zero_grad() lr_scheduler.step() #wait for all processes to finish accelerator.wait_for_everyone() # Checks if the accelerator has performed an optimization step behind the scenes if accelerator.sync_gradients: progress_bar.update(1) global_step += 1 if signal_recieved_time != 0: if time.time() - signal_recieved_time > 60: print("Signal received, saving state and exiting") atomic_save(save_path, accelerator) signal_recieved_time = 0 exit(0) else: exit(0) if accelerator.is_main_process: if global_step % args.checkpointing_steps == 0: atomic_save(save_path, accelerator) logger.info(f"Saved state to {save_path}") logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]} progress_bar.set_postfix(**logs) accelerator.log(logs, step=global_step) if global_step >= args.max_train_steps: break print("Step", step) accelerator.wait_for_everyone() if step == 0 or args.validation_prompt is not None and (step + 1) % args.validation_steps == 0: # Create pipeline pipe = ControlnetCogVideoXPipeline.from_pretrained( os.path.join(args.base_dir, args.pretrained_model_name_or_path), transformer=unwrap_model(transformer), text_encoder=unwrap_model(text_encoder), vae=unwrap_model(vae), scheduler=scheduler, torch_dtype=weight_dtype, ) print("Length of validation dataset: ", len(val_dataloader)) #create a pipeline per accelerator device (for faster inference) with torch.autocast(str(accelerator.device).replace(":0", ""), enabled=accelerator.mixed_precision == "fp16"): for batch in val_dataloader: frame = ((batch["blur_img"][0].permute(0,2,3,1).cpu().numpy() + 1)*127.5).astype(np.uint8) pipeline_args = { "prompt": "", "negative_prompt": "", "image": frame, "input_intervals": batch["input_intervals"][0:1], "output_intervals": batch["output_intervals"][0:1], "guidance_scale": args.guidance_scale, "use_dynamic_cfg": args.use_dynamic_cfg, "height": args.height, "width": args.width, "num_frames": args.max_num_frames, "num_inference_steps": args.num_inference_steps, } modified_filenames = [] filenames = batch['file_names'] for file in filenames: modified_filenames.append(os.path.splitext(file)[0] + ".mp4") num_frames = batch["num_frames"][0] #save the gt_video output if args.dataset not in ["outsidephotos"]: gt_video = batch["videos"][0].permute(0,2,3,1).cpu().numpy() gt_video = ((gt_video + 1) * 127.5)/255 gt_video = gt_video[0:num_frames] for file in modified_filenames: gt_file_name = os.path.join(args.output_dir, "gt", modified_filenames[0]) os.makedirs(os.path.dirname(gt_file_name), exist_ok=True) if args.dataset == "baist": bbox = batch["bbx"][0].cpu().numpy().astype(np.int32) gt_video = gt_video[:, bbox[1]:bbox[3], bbox[0]:bbox[2], :] gt_video = np.array([cv2.resize(frame, (160, 192)) for frame in gt_video]) #resize to 192x160 save_frames_as_pngs((gt_video*255).astype(np.uint8), gt_file_name.replace(".mp4", "").replace("gt", "gt_frames")) export_to_video(gt_video, gt_file_name, fps=20) if "high_fps_video" in batch: high_fps_video = batch["high_fps_video"][0].permute(0,2,3,1).cpu().numpy() high_fps_video = ((high_fps_video + 1) * 127.5)/255 gt_file_name = os.path.join(args.output_dir, "gt_highfps", modified_filenames[0]) #save the blurred image if args.dataset in ["full", "outsidephotos", "gopro2x", "baist"]: for file in modified_filenames: blurry_file_name = os.path.join(args.output_dir, "blurry", modified_filenames[0].replace(".mp4", ".png")) os.makedirs(os.path.dirname(blurry_file_name), exist_ok=True) if args.dataset == "baist": bbox = batch["bbx"][0].cpu().numpy().astype(np.int32) frame0 = frame[0][bbox[1]:bbox[3], bbox[0]:bbox[2], :] frame0 = cv2.resize(frame0, (160, 192)) #resize to 192x160 Image.fromarray(frame0).save(blurry_file_name) else: Image.fromarray(frame[0]).save(blurry_file_name) videos = log_validation( pipe=pipe, args=args, accelerator=accelerator, pipeline_args=pipeline_args ) #save the output video frames as pngs (uncompressed results) and mp4 (compressed results easily viewable) for i, video in enumerate(videos): video = video[0:num_frames] filename = os.path.join(args.output_dir, "deblurred", modified_filenames[0]) os.makedirs(os.path.dirname(filename), exist_ok=True) if args.dataset == "baist": bbox = batch["bbx"][0].cpu().numpy().astype(np.int32) video = video[:, bbox[1]:bbox[3], bbox[0]:bbox[2], :] video = np.array([cv2.resize(frame, (160, 192)) for frame in video]) #resize to 192x160 save_frames_as_pngs((video*255).astype(np.uint8), filename.replace(".mp4", "").replace("deblurred", "deblurred_frames")) export_to_video(video, filename, fps=20) accelerator.wait_for_everyone() if args.just_validate: exit(0) accelerator.wait_for_everyone() accelerator.end_training() signal_recieved_time = 0 def handle_signal(signum, frame): global signal_recieved_time signal_recieved_time = time.time() print(f"Signal {signum} received at {time.ctime()}") with open("/datasets/sai/gencam/cogvideox/interrupted.txt", "w") as f: f.write(f"Training was interrupted at {time.ctime()}") if __name__ == "__main__": args = get_args() print("Registering signal handler") #Register the signal handler (catch SIGUSR1) signal.signal(signal.SIGUSR1, handle_signal) main_thread = threading.Thread(target=main, args=(args,)) main_thread.start() while signal_recieved_time!= 0: time.sleep(1) #call main with args as a thread