o Ti;p@sddlZddlZddlZddlmZddlmZddlmZddl m Z m Z m Z m Z ddlZddlmZddlmZddlmZmZdd lmZdd lmZd d lmZd dlmZd dlmZm Z m!Z!erkddl"Z"ee#Z$GdddeZ%dS)N) defaultdict)futures)Path)AnyCallableOptionalUnion) Accelerator) get_logger)ProjectConfigurationset_seed)PyTorchModelHubMixin)is_wandb_available)DDPOStableDiffusionPipeline) DDPOConfig)PerPromptStatTrackergenerate_model_cardget_comet_experiment_urlcsReZdZdZddgZ d3dedeeje e e e gejfdege e e ffde d e ee e e ge ff d d Zd4d dZdedefddZddZdejdedejfddZddZddZddZd d!Zd"d#Zd$e ee ffd%d&Zd3d'e efd(d)Zd*d+Zfd,d-Z   d5d.e e d/e e d0ee e e dffd1d2Z!Z"S)6 DDPOTrainerah The DDPOTrainer uses Deep Diffusion Policy Optimization to optimise diffusion models. Note, this trainer is heavily inspired by the work here: https://github.com/kvablack/ddpo-pytorch As of now only Stable Diffusion based pipelines are supported Attributes: **config** (`DDPOConfig`) -- Configuration object for DDPOTrainer. Check the documentation of `PPOConfig` for more: details. **reward_function** (Callable[[torch.Tensor, tuple[str], tuple[Any]], torch.Tensor]) -- Reward function to be used: **prompt_function** (Callable[[], tuple[str, Any]]) -- Function to generate prompts to guide model **sd_pipeline** (`DDPOStableDiffusionPipeline`) -- Stable Diffusion pipeline to be used for training. **image_samples_hook** (Optional[Callable[[Any, Any, Any], Any]]) -- Hook to be called to log images trlddpoNconfigreward_functionprompt_function sd_pipelineimage_samples_hookc Cstdt|durtd||_||_||_||_td"i|jj}|jj r}t j t j |jj |j_ dt j |jj vr}ttddt |jj }t|dkr]td|jj tdd |D}t j |jj d|d |j_ |d d |_t|jj|jj|_td"|jj|jj||jj|jd |jj|_ |!\} } | st| |jduo|jd k} |j j"r|j j#|jj$| st%|&dn|&|jj'dt()d|t*|jj+dd||_,|j,j-d |j j. dddd|j jdkrt/j0} n|j jdkrt/j1} nt/j2} |j,j3j4|j j5| d|j,j6j4|j j5| d|j,j7j4|j j5| d|j,8} |j 9|j:|j ;|j<|jj=rJdt/j>j?j@_=|AtB| tsV| Cn| |_D|j,6|j,jE|jjFdurjdgn|jjFddd|j,jEjGdjH4|j j5d|_I|jJrtK|jL|jM|_N|j,jOp|j jO|_OtP|j,dr|j,jQr|j R| |jD\}|_Dttdd|C|_Sn |j R| |jD\|_S|_D|jjTrtUjV|jWd|_X|j rt()d |j |j Y|j t|j Zd!d d |_[dSd|_[dS)#Nz@DDPOTrainer is deprecated and will be removed in version 0.23.0.z8No image_samples_hook provided; no images will be logged checkpoint_cSsd|vS)NrxrrL/home/ubuntu/.local/lib/python3.10/site-packages/trl/trainer/ddpo_trainer.pyWsz&DDPOTrainer.__init__..rzNo checkpoints found in cSsg|] }t|ddqS)_)intsplit).0r!rrr" ]z(DDPOTrainer.__init__..r%r)log_withmixed_precisionproject_configgradient_accumulation_steps tensorboard)ddpo_trainer_config)r init_kwargs T)device_specificFTimestep)positiondisableleavedesc dynamic_ncolsfp16bf16)dtypept max_lengthreturn_tensorspadding truncationr?use_loracSs|jSN) requires_grad)prrr"r#s) max_workerszResuming from r$r)\warningswarnDeprecationWarning prompt_fn reward_fnrimage_samples_callbackr project_kwargs resume_fromospathnormpath expanduserbasenamelistfilterlistdirlen ValueErrorsortedjoin iterationr&sample_num_stepstrain_timestep_fractionnum_train_timestepsr r+r,!train_gradient_accumulation_stepsaccelerator_kwargs accelerator _config_checkis_main_process init_trackerstracker_project_namedictto_dicttracker_kwargsloggerinfor seedrset_progress_bar_configis_local_main_processtorchfloat16bfloat16float32vaetodevice text_encoderunetget_trainable_layersregister_save_state_pre_hook_save_model_hookregister_load_state_pre_hook_load_model_hook allow_tf32backendscudamatmul_setup_optimizer isinstance parameters optimizer tokenizernegative_promptsmodel_max_length input_idsneg_prompt_embedper_prompt_stat_trackingr$per_prompt_stat_tracking_buffer_size"per_prompt_stat_tracking_min_count stat_trackerautocasthasattrrDpreparetrainable_layersasync_reward_computationrThreadPoolExecutorrHexecutor load_stater' first_epoch)selfrrrrraccelerator_project_config checkpointscheckpoint_numbersis_okaymessageis_using_tensorboardinference_dtyperrxrrr"__init__;s              zDDPOTrainer.__init__Fc s~|s'g}|D]\}}}|||\}}|tj|jjd|fqt|Sjfdd|}fdd|D}t|S)Nrvcs j|SrE)rMr rrr"r#s z-DDPOTrainer.compute_rewards..cs.g|]\}}tj|jjd|fqSr)rp as_tensorresultrcrv)r(rewardreward_metadatarrr"r)sz/DDPOTrainer.compute_rewards..) rMappendrprrcrvrmapzip) rprompt_image_pairsis_asyncrewardsimagespromptsprompt_metadatarrrrr"compute_rewardss  zDDPOTrainer.compute_rewardsepoch global_stepcsjjjjjd\}fdddDj|jjd\}}t|D]\}}|||||gq)j durI ||j j dt |}j |}j j||||d|djjrj d }jjj|d d } j| |} n |||d } t | j jd j jj jd<d =dj \} t!jj"D]v} t j#| j jdfdd$Dt %fddt!| D} dD]}|t j&| j jddddf| f|<q'}fdd|D}t(|}fdd|D}jj)*+| |||}j j,s2t-dq|dkrK|jj.dkrKj j/rKj 0|S)a Perform a single step of training. Args: epoch (int): The current epoch. global_step (int): The current global step. Side Effects: - Model weights are updated - Logs the statistics to the accelerator trackers. - If `self.image_samples_callback` is not None, it will be called with the prompt_image_pairs, global_step, and the accelerator tracker. Returns: global_step (int): The updated global step. ) iterations batch_sizecs&i|]tfddDqS)csg|]}|qSrr)r(skrr"r)z/DDPOTrainer.step...)rpcat)r()samplesrr" s&z$DDPOTrainer.step..r)rN)rr reward_mean reward_stdstep prompt_idsT)skip_special_tokensg:0yE>r% advantages timestepsrcsi|] \}}||qSrrr(rv)permrr"r'csg|] }tjjjdqSr)rprandpermrcrvr(r$) num_timestepsrrr"r),r*z$DDPOTrainer.step..)rlatents next_latents log_probscs.g|]}|jdjjg|jddRqS)r%rN)reshapertrain_batch_sizeshape)r(rrrr"r)8s.csg|] }tt|qSr)rhr)r( row_values) original_keysrr"r)=rzsOptimization step should have been performed by this point. Please check calculated gradient accumulation settings.)1_generate_samplesrsample_num_batches_per_epochsample_batch_sizekeysrr enumerateextendrNrctrackersrprgathercpunumpylogmeanstdrrr batch_decoderupdaterr num_processes process_indexrurvrrangetrain_num_inner_epochsritemsstackarangevaluesrrxtrain_train_batched_samplessync_gradientsrZ save_freqre save_state)rrrprompt_image_datarrewards_metadatai image_datarrrtotal_batch_size inner_epochpermskeyoriginal_valuesreshaped_valuestransposed_valuessamples_batchedr)rrrrrr"rsz         & zDDPOTrainer.stepcCs(|L|jjr0|jt|gdt|gd|j}|d\}} ||jj | |}n |j|||j}|jj ||||jj |d} | j } Wdn1sSwYt ||jj |jj}t| |} |||jj| } dt| |d}tt| d|jjk}| ||fS)a Calculate the loss for a batch of an unpacked sample Args: latents (torch.Tensor): The latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, width] timesteps (torch.Tensor): The timesteps sampled from the diffusion model, shape: [batch_size] next_latents (torch.Tensor): The next latents sampled from the diffusion model, shape: [batch_size, num_channels_latents, height, width] log_probs (torch.Tensor): The log probabilities of the latents, shape: [batch_size] advantages (torch.Tensor): The advantages of the latents, shape: [batch_size] embeds (torch.Tensor): The embeddings of the prompts, shape: [2*batch_size or batch_size, ...] Note: the "or" is because if train_cfg is True, the expectation is that negative prompts are concatenated to the embeds Returns: loss (torch.Tensor), approx_kl (torch.Tensor), clipfrac (torch.Tensor) (all of these are of shape (1,)) r)eta prev_sampleNg??)rr train_cfgrrxrprsamplechunksample_guidance_scalescheduler_step sample_etarclamptrain_adv_clip_maxexplosstrain_clip_rangerabsfloat)rrrrrrembeds noise_prednoise_pred_uncondnoise_pred_textscheduler_step_outputlog_probratior approx_klclipfracrrr"calculate_lossLsN   zDDPOTrainer.calculate_lossr clip_rangercCs8| |}| t|d|d|}tt||S)Nr)rprrmaximum)rrrrunclipped_loss clipped_lossrrr"rs zDDPOTrainer.losscCsL|jjr ddl}|jj}ntjj}|||jj|jj|jj f|jj |jj dS)Nr)lrbetas weight_decayeps) rtrain_use_8bit_adam bitsandbytesoptim AdamW8bitrpAdamWtrain_learning_ratetrain_adam_beta1train_adam_beta2train_adam_weight_decaytrain_adam_epsilon)rtrainable_layers_parametersr optimizer_clsrrr"rs zDDPOTrainer._setup_optimizercCs|j||||dSrE)rsave_checkpointpop)rmodelsweights output_dirrrr"r{s zDDPOTrainer._save_model_hookcCs|j|||dSrE)rload_checkpointr+)rr, input_dirrrr"r}s zDDPOTrainer._load_model_hookc sdg}g}jjj|dd}t|D]}tfddt|D\}}jj|dddjjjdj j j } j | d} "j| |jjjjjjdd } | j} | j} | j}Wd n1slwYtj| dd } tj|dd }jjj|d}|| | || d d d d f| d d dd f||d || ||gq||fS)a4 Generate samples from the model Args: iterations (int): Number of iterations to generate samples for batch_size (int): Batch size to use for sampling Returns: samples (list[dict[str, torch.Tensor]]), prompt_image_pairs (list[list[Any]]) rcsg|]}qSr)rLrrrr"r)rz1DDPOTrainer._generate_samples..r>r?Tr@r) prompt_embedsnegative_prompt_embedsnum_inference_stepsguidance_scaler output_typeN)dimr%)rr1rrrrr2)rrxevalrrepeatrrrrrrurcrvrwrrr^rrrrrrpr schedulerrr)rrrrrsample_neg_prompt_embedsr$rrrr1 sd_outputrrrrrrr"rsX       zDDPOTrainer._generate_samplesc Cstt}t|D]\}}|jjrt|d|dg}n|d}t|jD]} |j |j j u| |ddd| f|ddd| f|ddd| f|ddd| f|d|\} } } |d | |d | |d | |j | |j jr|j t|jts|jn|j|jj|j|jWdn1swY|j jrd d |D}|j j|dd}|||d|j j||d|d7}tt}q%q|S)a Train on a batch of samples. Main training segment Args: inner_epoch (int): The current inner epoch epoch (int): The current epoch global_step (int): The current global step batched_samples (list[dict[str, torch.Tensor]]): The batched samples to train on Side Effects: - Model weights are updated - Logs the statistics to the accelerator trackers. Returns: global_step (int): The updated global step r2r1rNrrrrrrrcSs"i|] \}}|tt|qSr)rprrrrrr"r)s"z6DDPOTrainer._train_batched_samples..r) reduction)rrrr)rrVrrrrprrr`rc accumulaterrxrrbackwardrclip_grad_norm_rrrtrain_max_grad_normrr zero_gradrreducerr) rrrrbatched_samplesrl_irr jrrrrrr"rsN      "z"DDPOTrainer._train_batched_samplesreturncCs|jj|jj|jj}|jj|jj|jj}|jj|jjks/dd|jjd|jjdfS|jj|jjdksHdd|jjd|jjdfS||dksYdd|d|dfSd S) NFzSample batch size (z9) must be greater than or equal to the train batch size ()rz-) must be divisible by the train batch size (zNumber of samples per epoch (z3) must be divisible by the total train batch size ()Tr=)rrrcrrrra)rsamples_per_epochtotal_train_batch_sizerrr"rd1s* zDDPOTrainer._config_checkepochscCs6d}|dur |jj}t|j|D]}|||}qdS)z> Train the model for a given number of epochs rN)r num_epochsrrr)rrJrrrrr"rLs zDDPOTrainer.traincCs|j||dSrE)rsave_pretrainedcreate_model_card)rsave_directoryrrr"_save_pretrainedVs  zDDPOTrainer._save_pretrainedcsL|jjdurt|jjj}n |jjdd}|j|dt||dS)N/r%) model_name) args hub_model_idrr.namer'rMsuper_save_checkpoint)rmodeltrialrQ __class__rr"rV[s  zDDPOTrainer._save_checkpointrQ dataset_nametagsc Cs|sdSt|jjdrtj|jjjs|jjj}nd}|dur&t}n t |t r/|h}nt|}t|jjdr?| d| |j td}t|||j||tr]tjdur]tjjndtd|ddd }|tj|jjd dS) a Creates a draft of a model card using the information available to the `Trainer`. Args: model_name (`str` or `None`, *optional*, defaults to `None`): Name of the model. dataset_name (`str` or `None`, *optional*, defaults to `None`): Name of the dataset used for training. tags (`str`, `list[str]` or `None`, *optional*, defaults to `None`): Tags to be associated with the model card. N _name_or_pathunsloth_versionunslotha @inproceedings{black2024training, title = {{Training Diffusion Models with Reinforcement Learning}}, author = {Kevin Black and Michael Janner and Yilun Du and Ilya Kostrikov and Sergey Levine}, year = 2024, booktitle = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024}, publisher = {OpenReview.net}, url = {https://openreview.net/forum?id=YCWjhGrJFD}, }DDPOz5Training Diffusion Models with Reinforcement Learningz 2305.13301) base_modelrQrSr[r\ wandb_url comet_url trainer_nametrainer_citation paper_titlepaper_idz README.md)is_world_process_zerorrWrrQrRisdirr]setrstraddr _tag_namestextwrapdedentrrSrwandbrunurlrsaver\rRr.)rrQr[r\racitation model_cardrrr"rMcs8       zDDPOTrainer.create_model_cardrE)F)NNN)#__name__ __module__ __qualname____doc__rmrrrpTensortuplerkrrrrrr&rrr rrr{r}rrboolrdrrOrVrrVrM __classcell__rrrYr"r*sZ  kD >=  r)&rQrnrI collectionsr concurrentrpathlibrtypingrrrrrp accelerater accelerate.loggingr accelerate.utilsr r huggingface_hubr transformersrr,r ddpo_configrutilsrrrrprvrkrrrrr"s(