
"""simple docstring"""


import re
import time
from typing import Optional

import IPython.display as disp

from ..trainer_callback import TrainerCallback
from ..trainer_utils import IntervalStrategy, has_length



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				int:







       """simple docstring"""
       UpperCamelCase       			=		int(__A					)
       UpperCamelCase			,							UpperCamelCase			,							UpperCamelCase       			=		t // 36_00, (t // 60) % 60, t % 60
       return F"{h}:{m:02d}:{s:02d}" if h != 0 else F"{m:02d}:{s:02d}"



def 						lowerCamelCase__      (				UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_		,   UpperCAmelCase_=3_00					)->				List[str]:







       """simple docstring"""
       # docstyle-ignore
       return F"\n    <div>\n      {prefix}\n      <progress value=\'{value}\' max=\'{total}\' style=\'width:{width}px; height:20px; vertical-align: middle;\'></progress>\n      {label}\n    </div>\n    "



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				Tuple:







       """simple docstring"""
       UpperCamelCase       			=		"<table border=\"1\" class=\"dataframe\">\n"
       html_code += """  <thead>\n <tr style="text-align: left;">\n"""
       for i in items[0]:
              html_code += F"      <th>{i}</th>\n"
       html_code += "    </tr>\n  </thead>\n  <tbody>\n"
       for line in items[1:]:
              html_code += "    <tr>\n"
              for elt in line:
                     UpperCamelCase       			=		F"{elt:.6f}" if isinstance(__A		,   __A					) else str(__A					)
                     html_code += F"      <td>{elt}</td>\n"
              html_code += "    </tr>\n"
       html_code += "  </tbody>\n</table><p>"
       return html_code



class 	__a	:
       UpperCamelCase_		:       Tuple				=      5
       UpperCamelCase_		:       Tuple				=      0.2






       def __init__(	self			:				List[str]							,   UpperCAmelCase_			:				Optional[int]							,   UpperCAmelCase_			:				List[Any] = None							,   UpperCAmelCase_			:				Any = True							,   UpperCAmelCase_			:				Optional[int] = None							,   UpperCAmelCase_			:				Optional[int] = 300							,   )->	str:







              """simple docstring"""
              UpperCamelCase       			=		total
              UpperCamelCase       			=		"" if prefix is None else prefix
              UpperCamelCase       			=		leave
              UpperCamelCase       			=		parent
              UpperCamelCase       			=		width
              UpperCamelCase       			=		None
              UpperCamelCase       			=		None
              UpperCamelCase       			=		None






       def 						_SCREAMING_SNAKE_CASE							(	self			:				int							,   UpperCAmelCase_			:				str							,   UpperCAmelCase_			:				List[Any] = False							,   UpperCAmelCase_			:				int = None      )->	Dict:







              """simple docstring"""
              UpperCamelCase       			=		value
              if comment is not None:
                     UpperCamelCase       			=		comment
              if self.last_value is None:
                     UpperCamelCase       			=		UpperCamelCase       			=		time.time()
                     UpperCamelCase       			=		UpperCamelCase       			=		value
                     UpperCamelCase       			=		UpperCamelCase       			=		None
                     UpperCamelCase       			=		self.warmup
                     UpperCamelCase       			=		1
                     self.update_bar(lowercase__      )
              elif value <= self.last_value and not force_update:
                     return
              elif force_update or self.first_calls > 0 or value >= min(self.last_value + self.wait_for							,   self.total      ):
                     if self.first_calls > 0:
                            self.first_calls -= 1
                     UpperCamelCase       			=		time.time()
                     UpperCamelCase       			=		current_time - self.start_time
                     # We could have value = self.start_value if the update is called twixe with the same start value.
                     if value > self.start_value:
                            UpperCamelCase       			=		self.elapsed_time / (value - self.start_value)
                     else:
                            UpperCamelCase       			=		None
                     if value >= self.total:
                            UpperCamelCase       			=		self.total
                            UpperCamelCase       			=		None
                            if not self.leave:
                                   self.close()
                     elif self.average_time_per_item is not None:
                            UpperCamelCase       			=		self.average_time_per_item * (self.total - value)
                     self.update_bar(lowercase__      )
                     UpperCamelCase       			=		value
                     UpperCamelCase       			=		current_time
                     if self.average_time_per_item is None:
                            UpperCamelCase       			=		1
                     else:
                            UpperCamelCase       			=		max(int(self.update_every / self.average_time_per_item      )							,   1      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any							,   UpperCAmelCase_			:				Optional[Any]							,   UpperCAmelCase_			:				int=None      )->	Optional[Any]:







              """simple docstring"""
              UpperCamelCase       			=		" " * (len(str(self.total      )      ) - len(str(lowercase__      )      )) + str(lowercase__      )
              if self.elapsed_time is None:
                     UpperCamelCase       			=		f"[{spaced_value}/{self.total} : < :"
              elif self.predicted_remaining is None:
                     UpperCamelCase       			=		f"[{spaced_value}/{self.total} {format_time(self.elapsed_time      )}"
              else:
                     UpperCamelCase       			=		(
                         f"[{spaced_value}/{self.total} {format_time(self.elapsed_time      )} <"
                         f" {format_time(self.predicted_remaining      )}"
                     )
                     self.label += f", {1/self.average_time_per_item:.2f} it/s"
              self.label += "]" if self.comment is None or len(self.comment      ) == 0 else f", {self.comment}]"
              self.display()






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	Dict:







              """simple docstring"""
              UpperCamelCase       			=		html_progress_bar(self.value							,   self.total							,   self.prefix							,   self.label							,   self.width      )
              if self.parent is not None:
                     # If this is a child bar, the parent will take care of the display.
                     self.parent.display()
                     return
              if self.output is None:
                     UpperCamelCase       			=		disp.display(disp.HTML(self.html_code      )							,   display_id=lowercase__      )
              else:
                     self.output.update(disp.HTML(self.html_code      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[Any]      )->	List[str]:







              """simple docstring"""



              if self.parent is None and self.output is not None:
                     self.output.update(disp.HTML(""      )      )





class 	__a	(   _lowerCAmelCase						):






       def __init__(	self			:				int							,   UpperCAmelCase_			:				Dict							,   UpperCAmelCase_			:				str=None      )->	int:







              """simple docstring"""
              super().__init__(lowercase__      )
              UpperCamelCase       			=		None if column_names is None else [column_names]
              UpperCamelCase       			=		None






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	int:







              """simple docstring"""
              UpperCamelCase       			=		html_progress_bar(self.value							,   self.total							,   self.prefix							,   self.label							,   self.width      )
              if self.inner_table is not None:
                     self.html_code += text_to_html_table(self.inner_table      )
              if self.child_bar is not None:
                     self.html_code += self.child_bar.html_code
              if self.output is None:
                     UpperCamelCase       			=		disp.display(disp.HTML(self.html_code      )							,   display_id=lowercase__      )
              else:
                     self.output.update(disp.HTML(self.html_code      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Union[str, Any]							,   UpperCAmelCase_			:				int      )->	Optional[Any]:







              """simple docstring"""



              if self.inner_table is None:
                     UpperCamelCase       			=		[list(values.keys()      ), list(values.values()      )]
              else:
                     UpperCamelCase       			=		self.inner_table[0]
                     if len(self.inner_table      ) == 1:
                            # We give a chance to update the column names at the first iteration
                            for key in values.keys():
                                   if key not in columns:
                                          columns.append(lowercase__      )
                            UpperCamelCase       			=		columns
                     self.inner_table.append([values[c] for c in columns]      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[Any]							,   UpperCAmelCase_			:				Tuple							,   UpperCAmelCase_			:				Tuple=None							,   UpperCAmelCase_			:				List[Any]=300      )->	Union[str, Any]:







              """simple docstring"""
              UpperCamelCase       			=		NotebookProgressBar(lowercase__							,   prefix=lowercase__							,   parent=self							,   width=lowercase__      )
              return self.child_bar






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[str]      )->	List[Any]:







              """simple docstring"""
              UpperCamelCase       			=		None
              self.display()





class 	__a	(   _lowerCAmelCase						):






       def __init__(	self			:				List[Any]      )->	str:







              """simple docstring"""
              UpperCamelCase       			=		None
              UpperCamelCase       			=		None
              UpperCamelCase       			=		False






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]							,   UpperCAmelCase_			:				Any							,   UpperCAmelCase_			:				str							,   UpperCAmelCase_			:				Optional[int]							,   **UpperCAmelCase_			:				List[str]      )->	Union[str, Any]:







              """simple docstring"""
              UpperCamelCase       			=		"Epoch" if args.evaluation_strategy == IntervalStrategy.EPOCH else "Step"
              UpperCamelCase       			=		0
              UpperCamelCase       			=		0
              UpperCamelCase       			=		[self.first_column] + ["Training Loss"]
              if args.evaluation_strategy != IntervalStrategy.NO:
                     column_names.append("Validation Loss"      )
              UpperCamelCase       			=		NotebookTrainingTracker(state.max_steps							,   lowercase__      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[str]							,   UpperCAmelCase_			:				Optional[int]							,   UpperCAmelCase_			:				Union[str, Any]							,   UpperCAmelCase_			:				List[str]							,   **UpperCAmelCase_			:				Optional[int]      )->	List[str]:







              """simple docstring"""
              UpperCamelCase       			=		int(state.epoch      ) if int(state.epoch      ) == state.epoch else f"{state.epoch:.2f}"
              self.training_tracker.update(
                  state.global_step + 1							,   comment=f"Epoch {epoch}/{state.num_train_epochs}"							,   force_update=self._force_next_update							,   )
              UpperCamelCase       			=		False






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]							,   UpperCAmelCase_			:				Union[str, Any]							,   UpperCAmelCase_			:				str							,   UpperCAmelCase_			:				Tuple							,   UpperCAmelCase_			:				Union[str, Any]=None							,   **UpperCAmelCase_			:				Optional[Any]      )->	Tuple:







              """simple docstring"""



              if not has_length(lowercase__      ):
                     return
              if self.prediction_bar is None:
                     if self.training_tracker is not None:
                            UpperCamelCase       			=		self.training_tracker.add_child(len(lowercase__      )      )
                     else:
                            UpperCamelCase       			=		NotebookProgressBar(len(lowercase__      )      )
                     self.prediction_bar.update(1      )
              else:
                     self.prediction_bar.update(self.prediction_bar.value + 1      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]							,   UpperCAmelCase_			:				str							,   UpperCAmelCase_			:				Optional[int]							,   UpperCAmelCase_			:				List[Any]							,   **UpperCAmelCase_			:				str      )->	List[str]:







              """simple docstring"""



              if self.prediction_bar is not None:
                     self.prediction_bar.close()
              UpperCamelCase       			=		None






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]							,   UpperCAmelCase_			:				Optional[Any]							,   UpperCAmelCase_			:				Optional[Any]							,   UpperCAmelCase_			:				Optional[Any]							,   UpperCAmelCase_			:				str=None							,   **UpperCAmelCase_			:				Union[str, Any]      )->	int:







              """simple docstring"""



              if args.evaluation_strategy == IntervalStrategy.NO and "loss" in logs:
                     UpperCamelCase       			=		{"Training Loss": logs["loss"]}
                     # First column is necessarily Step sine we're not in epoch eval strategy
                     UpperCamelCase       			=		state.global_step
                     self.training_tracker.write_line(lowercase__      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Any							,   UpperCAmelCase_			:				int							,   UpperCAmelCase_			:				Tuple							,   UpperCAmelCase_			:				Optional[Any]							,   UpperCAmelCase_			:				Tuple=None							,   **UpperCAmelCase_			:				List[Any]      )->	Union[str, Any]:







              """simple docstring"""



              if self.training_tracker is not None:
                     UpperCamelCase       			=		{"Training Loss": "No log", "Validation Loss": "No log"}
                     for log in reversed(state.log_history      ):
                            if "loss" in log:
                                   UpperCamelCase       			=		log["loss"]
                                   break

                     if self.first_column == "Epoch":
                            UpperCamelCase       			=		int(state.epoch      )
                     else:
                            UpperCamelCase       			=		state.global_step
                     UpperCamelCase       			=		"eval"
                     for k in metrics:
                            if k.endswith("_loss"      ):
                                   UpperCamelCase       			=		re.sub(r"\_loss$"							,   ""							,   lowercase__      )
                     UpperCamelCase       			=		metrics.pop("total_flos"							,   lowercase__      )
                     UpperCamelCase       			=		metrics.pop("epoch"							,   lowercase__      )
                     UpperCamelCase       			=		metrics.pop(f"{metric_key_prefix}_runtime"							,   lowercase__      )
                     UpperCamelCase       			=		metrics.pop(f"{metric_key_prefix}_samples_per_second"							,   lowercase__      )
                     UpperCamelCase       			=		metrics.pop(f"{metric_key_prefix}_steps_per_second"							,   lowercase__      )
                     UpperCamelCase       			=		metrics.pop(f"{metric_key_prefix}_jit_compilation_time"							,   lowercase__      )
                     for k, v in metrics.items():
                            if k == f"{metric_key_prefix}_loss":
                                   UpperCamelCase       			=		v
                            else:
                                   UpperCamelCase       			=		k.split("_"      )
                                   UpperCamelCase       			=		" ".join([part.capitalize() for part in splits[1:]]      )
                                   UpperCamelCase       			=		v
                     self.training_tracker.write_line(lowercase__      )
                     self.training_tracker.remove_child()
                     UpperCamelCase       			=		None
                     # Evaluation takes a long time so we should force the next update.
                     UpperCamelCase       			=		True






       def 						_SCREAMING_SNAKE_CASE							(	self			:				List[Any]							,   UpperCAmelCase_			:				Dict							,   UpperCAmelCase_			:				Tuple							,   UpperCAmelCase_			:				int							,   **UpperCAmelCase_			:				int      )->	Dict:







              """simple docstring"""
              self.training_tracker.update(
                  state.global_step							,   comment=f"Epoch {int(state.epoch      )}/{state.num_train_epochs}"							,   force_update=lowercase__      )
              UpperCamelCase       			=		None








import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





@slow
@require_torch_gpu
class        lowerCAmelCase__  (					unittest.TestCase			):



       '''simple docstring'''







       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''     )
             __A    =sd_pipe.to(lowercase__     )
             sd_pipe.set_progress_bar_config(disable=lowercase__     )

             sd_pipe.set_scheduler('''sample_euler'''     )

             __A    ='''A painting of a squirrel eating a burger'''
             __A    =torch.manual_seed(0     )
             __A    =sd_pipe([prompt] ,   generator=lowercase__ ,   guidance_scale=9.0 ,   num_inference_steps=2_0 ,   output_type='''np'''     )

             __A    =output.images

             __A    =image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             __A    =np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]     )

             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2

       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''     )
             __A    =sd_pipe.to(lowercase__     )
             sd_pipe.set_progress_bar_config(disable=lowercase__     )

             sd_pipe.set_scheduler('''sample_euler'''     )

             __A    ='''A painting of a squirrel eating a burger'''
             __A    =torch.manual_seed(0     )
             __A    =sd_pipe([prompt] ,   generator=lowercase__ ,   guidance_scale=9.0 ,   num_inference_steps=2_0 ,   output_type='''np'''     )

             __A    =output.images

             __A    =image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             __A    =np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]     )

             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 5E-1


       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             __A    =StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''     )
             __A    =sd_pipe.to(lowercase__     )
             sd_pipe.set_progress_bar_config(disable=lowercase__     )

             sd_pipe.set_scheduler('''sample_dpmpp_2m'''     )

             __A    ='''A painting of a squirrel eating a burger'''
             __A    =torch.manual_seed(0     )
             __A    =sd_pipe(
                 [prompt] ,   generator=lowercase__ ,   guidance_scale=7.5 ,   num_inference_steps=1_5 ,   output_type='''np''' ,   use_karras_sigmas=lowercase__ ,   )

             __A    =output.images

             __A    =image[0, -3:, -3:, -1]

             assert image.shape == (1, 5_1_2, 5_1_2, 3)
             __A    =np.array(
                 [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]     )

             assert np.abs(image_slice.flatten() - expected_slice     ).max() < 1E-2


