





'''simple docstring'''






import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()




class 				UpperCAmelCase  (    unittest.TestCase     ):




      def        lowercase__					(    self    :				Optional[int]						)       ->   List[str]:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      @property
      def        lowercase__					(    self    :				int						)       ->   Optional[int]:
            _lowerCAmelCase							  =							1
            _lowerCAmelCase							  =							3
            _lowerCAmelCase							  =							(32, 32)

            _lowerCAmelCase							  =							floats_tensor((batch_size, num_channels) + sizes    ,    rng=random.Random(0						)						).to(__snake_case						)
            return image




      @property
      def        lowercase__					(    self    :				Tuple						)       ->   str:
            torch.manual_seed(0						)
            _lowerCAmelCase							  =							UNetaDConditionModel(
                block_out_channels=(32, 64)    ,    layers_per_block=2    ,    sample_size=32    ,    in_channels=4    ,    out_channels=4    ,    down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")    ,    up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")    ,    cross_attention_dim=32    ,    )
            return model




      @property
      def        lowercase__					(    self    :				Dict						)       ->   Optional[Any]:
            torch.manual_seed(0						)
            _lowerCAmelCase							  =							AutoencoderKL(
                block_out_channels=[32, 64]    ,    in_channels=3    ,    out_channels=3    ,    down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]    ,    up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]    ,    latent_channels=4    ,    )
            return model




      @property
      def        lowercase__					(    self    :				str						)       ->   str:
            torch.manual_seed(0						)
            _lowerCAmelCase							  =							CLIPTextConfig(
                bos_token_id=0    ,    eos_token_id=2    ,    hidden_size=32    ,    intermediate_size=37    ,    layer_norm_eps=1E-0_5    ,    num_attention_heads=4    ,    num_hidden_layers=5    ,    pad_token_id=1    ,    vocab_size=10_00    ,    )
            return CLIPTextModel(__snake_case						)




      @property
      def        lowercase__					(    self    :				List[Any]						)       ->   List[Any]:
            def extract(*__snake_case    :				str    ,    **__snake_case    :				Optional[Any]						):




                  class 				UpperCAmelCase  :




                        def __init__(    self    :				str						)       ->   int:
                              _lowerCAmelCase							  =							torch.ones([0]						)




                        def        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				Optional[Any]						)       ->   Optional[Any]:
                              self.pixel_values.to(__snake_case						)
                              return self




                  return Out()

            return extract




      def        lowercase__					(    self    :				int						)       ->   Tuple:
            _lowerCAmelCase							  =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase							  =							self.dummy_cond_unet
            _lowerCAmelCase							  =							DDIMScheduler(
                beta_start=0.0_00_85    ,    beta_end=0.0_12    ,    beta_schedule="""scaled_linear"""    ,    clip_sample=__snake_case    ,    set_alpha_to_one=__snake_case    ,    )

            _lowerCAmelCase							  =							self.dummy_vae
            _lowerCAmelCase							  =							self.dummy_text_encoder
            _lowerCAmelCase							  =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase							  =							StableDiffusionPipeline(
                unet=__snake_case    ,    scheduler=__snake_case    ,    vae=__snake_case    ,    text_encoder=__snake_case    ,    tokenizer=__snake_case    ,    safety_checker=__snake_case    ,    feature_extractor=self.dummy_extractor    ,    )
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""A painting of a squirrel eating a burger"""

            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe([prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""						)
            _lowerCAmelCase							  =							output.images

            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""    ,    return_dict=__snake_case    ,    )[0]

            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _lowerCAmelCase							  =							np.array([0.57_53, 0.61_14, 0.50_01, 0.50_34, 0.54_70, 0.47_29, 0.49_71, 0.48_67, 0.48_67]						)

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




      def        lowercase__					(    self    :				str						)       ->   str:
            _lowerCAmelCase							  =							"""cpu"""  # ensure determinism for the device-dependent torch.Generator
            _lowerCAmelCase							  =							self.dummy_cond_unet
            _lowerCAmelCase							  =							PNDMScheduler(skip_prk_steps=__snake_case						)
            _lowerCAmelCase							  =							self.dummy_vae
            _lowerCAmelCase							  =							self.dummy_text_encoder
            _lowerCAmelCase							  =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase							  =							StableDiffusionPipeline(
                unet=__snake_case    ,    scheduler=__snake_case    ,    vae=__snake_case    ,    text_encoder=__snake_case    ,    tokenizer=__snake_case    ,    safety_checker=__snake_case    ,    feature_extractor=self.dummy_extractor    ,    )
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe([prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""						)

            _lowerCAmelCase							  =							output.images

            _lowerCAmelCase							  =							torch.Generator(device=__snake_case						).manual_seed(0						)
            _lowerCAmelCase							  =							sd_pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=6.0    ,    num_inference_steps=2    ,    output_type="""np"""    ,    return_dict=__snake_case    ,    )[0]

            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							image_from_tuple[0, -3:, -3:, -1]

            assert image.shape == (1, 64, 64, 3)
            _lowerCAmelCase							  =							np.array([0.51_22, 0.57_12, 0.48_25, 0.50_53, 0.56_46, 0.47_69, 0.51_79, 0.48_94, 0.49_94]						)

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




      def        lowercase__					(    self    :				List[str]						)       ->   int:
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained(
                """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""    ,    safety_checker=__snake_case						)
            assert isinstance(__snake_case    ,    __snake_case						)
            assert isinstance(pipe.scheduler    ,    __snake_case						)
            assert pipe.safety_checker is None

            _lowerCAmelCase							  =							pipe("""example prompt"""    ,    num_inference_steps=2						).images[0]
            assert image is not None

            # check that there's no error when saving a pipeline with one of the models being None
            with tempfile.TemporaryDirectory() as tmpdirname:
                  pipe.save_pretrained(__snake_case						)
                  _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained(__snake_case						)

            # sanity check that the pipeline still works
            assert pipe.safety_checker is None
            _lowerCAmelCase							  =							pipe("""example prompt"""    ,    num_inference_steps=2						).images[0]
            assert image is not None




      @unittest.skipIf(torch_device != """cuda"""    ,    """This test requires a GPU"""						)
      def        lowercase__					(    self    :				Union[str, Any]						)       ->   List[Any]:
            _lowerCAmelCase							  =							self.dummy_cond_unet
            _lowerCAmelCase							  =							PNDMScheduler(skip_prk_steps=__snake_case						)
            _lowerCAmelCase							  =							self.dummy_vae
            _lowerCAmelCase							  =							self.dummy_text_encoder
            _lowerCAmelCase							  =							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""						)

            # put models in fp16
            _lowerCAmelCase							  =							unet.half()
            _lowerCAmelCase							  =							vae.half()
            _lowerCAmelCase							  =							bert.half()

            # make sure here that pndm scheduler skips prk
            _lowerCAmelCase							  =							StableDiffusionPipeline(
                unet=__snake_case    ,    scheduler=__snake_case    ,    vae=__snake_case    ,    text_encoder=__snake_case    ,    tokenizer=__snake_case    ,    safety_checker=__snake_case    ,    feature_extractor=self.dummy_extractor    ,    )
            _lowerCAmelCase							  =							sd_pipe.to(__snake_case						)
            sd_pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""A painting of a squirrel eating a burger"""
            _lowerCAmelCase							  =							sd_pipe([prompt]    ,    num_inference_steps=2    ,    output_type="""np"""						).images

            assert image.shape == (1, 64, 64, 3)




@nightly
@require_torch_gpu
class 				UpperCAmelCase  (    unittest.TestCase     ):




      def        lowercase__					(    self    :				List[str]						)       ->   List[str]:
            # clean up the VRAM after each test
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def        lowercase__					(    self    :				Optional[int]						)       ->   Tuple:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""a photo of a cat"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": ["""sunglasses"""],
                """reverse_editing_direction""": [False],
                """edit_warmup_steps""": 10,
                """edit_guidance_scale""": 6,
                """edit_threshold""": 0.95,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							3
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.34_67_31_13,
                0.38_49_27_33,
                0.37_59_73_52,
                0.34_08_63_35,
                0.35_65_07_48,
                0.35_57_92_05,
                0.3_38_47_63,
                0.34_34_02_36,
                0.3_57_32_71,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.41_88_78_26,
                0.37_72_87_66,
                0.30_13_82_72,
                0.41_41_63_35,
                0.41_66_49_85,
                0.36_28_33_92,
                0.36_19_12_46,
                0.43_36_44_65,
                0.43_00_17_32,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




      def        lowercase__					(    self    :				Optional[Any]						)       ->   Union[str, Any]:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""an image of a crowded boulevard, realistic, 4k"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": """crowd, crowded, people""",
                """reverse_editing_direction""": True,
                """edit_warmup_steps""": 10,
                """edit_guidance_scale""": 8.3,
                """edit_threshold""": 0.9,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							9
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.43_49_79_98,
                0.91_81_40_65,
                0.7_54_07_39,
                0.55_58_02_05,
                0.8_46_72_65,
                0.5_38_96_91,
                0.62_57_45_06,
                0.58_89_77_63,
                0.50_92_67_57,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.3_08_97_19,
                0.30_50_01_44,
                0.29_01_60_42,
                0.30_63_09_64,
                0.32_56_87,
                0.29_41_92_25,
                0.2_90_80_91,
                0.28_72_35_98,
                0.27_69_62_94,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




      def        lowercase__					(    self    :				int						)       ->   str:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""a castle next to a river"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": ["""boat on a river, boat""", """monet, impression, sunrise"""],
                """reverse_editing_direction""": False,
                """edit_warmup_steps""": [15, 18],
                """edit_guidance_scale""": 6,
                """edit_threshold""": [0.9, 0.8],
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							48
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.75_16_35_55,
                0.76_03_71_45,
                0.6_17_85,
                0.9_18_96_73,
                0.8_62_77_01,
                0.85_18_96_94,
                0.8_51_28_13,
                0.87_01_20_76,
                0.8_31_28_57,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.73_55_33_65,
                0.7_53_72_71,
                0.74_34_19_05,
                0.66_48_03_56,
                0.6_47_29_25,
                0.63_03_94_16,
                0.64_81_29_05,
                0.6_74_97_17,
                0.6_51_71_02,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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




      def        lowercase__					(    self    :				List[Any]						)       ->   Dict:
            _lowerCAmelCase							  =							"""cuda"""
            _lowerCAmelCase							  =							StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""    ,    torch_dtype=torch.floataa						)
            _lowerCAmelCase							  =							pipe.to(__snake_case						)
            pipe.set_progress_bar_config(disable=__snake_case						)

            _lowerCAmelCase							  =							"""a photo of a cat"""
            _lowerCAmelCase							  =							{
                """editing_prompt""": ["""sunglasses"""],
                """reverse_editing_direction""": [False],
                """edit_warmup_steps""": 10,
                """edit_guidance_scale""": 6,
                """edit_threshold""": 0.95,
                """edit_momentum_scale""": 0.5,
                """edit_mom_beta""": 0.6,
            }

            _lowerCAmelCase							  =							3
            _lowerCAmelCase							  =							7

            # no sega enabled
            _lowerCAmelCase							  =							torch.Generator(__snake_case						)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.34_88_76_95,
                0.3_87_69_53,
                0.3_75,
                0.34_42_38_28,
                0.3_58_15_43,
                0.35_71_77_73,
                0.3_38_37_89,
                0.34_57_03_12,
                0.35_93_75,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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

            # with sega enabled
            # generator = torch.manual_seed(seed)
            generator.manual_seed(__snake_case						)
            _lowerCAmelCase							  =							pipe(
                [prompt]    ,    generator=__snake_case    ,    guidance_scale=__snake_case    ,    num_inference_steps=50    ,    output_type="""np"""    ,    width=5_12    ,    height=5_12    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							output.images
            _lowerCAmelCase							  =							image[0, -3:, -3:, -1]
            _lowerCAmelCase							  =							[
                0.42_28_51_56,
                0.36_91_40_62,
                0.29_07_71_48,
                0.42_04_10_16,
                0.41_91_89_45,
                0.35_49_80_47,
                0.3_61_81_64,
                0.4_42_38_28,
                0.43_11_52_34,
            ]

            assert image.shape == (1, 5_12, 5_12, 3)

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













'''simple docstring'''






import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
      import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


A__       :     Optional[Any]          =get_logger(__name__)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = None							,			):





      """simple docstring"""


      for i, state in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Model weights saved in {output_model_file}"						)
      # Optimizer states
      for i, opt in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							opt.state_dict()
            _lowerCAmelCase							  =							f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Optimizer state saved in {output_optimizer_file}"						)
      # Scheduler states
      for i, scheduler in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							scheduler.state_dict()
            _lowerCAmelCase							  =							f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Scheduler state saved in {output_scheduler_file}"						)
      # GradScaler state
      if scaler is not None:
            _lowerCAmelCase							  =							scaler.state_dict()
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            torch.save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Gradient scaler state saved in {output_scaler_file}"						)
      # Random number generator states
      _lowerCAmelCase							  =							{}
      _lowerCAmelCase							  =							f"{RNG_STATE_NAME}_{process_index}.pkl"
      _lowerCAmelCase							  =							random.getstate()
      _lowerCAmelCase							  =							np.random.get_state()
      _lowerCAmelCase							  =							torch.get_rng_state()
      if is_xpu_available():
            _lowerCAmelCase							  =							torch.xpu.get_rng_state_all()
      else:
            _lowerCAmelCase							  =							torch.cuda.get_rng_state_all()
      if is_tpu_available():
            _lowerCAmelCase							  =							xm.get_rng_state()
      _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
      torch.save(lowerCAmelCase							,			lowerCAmelCase						)
      logger.info(f"Random states saved in {output_states_file}"						)
      return output_dir
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase=None							,			lowerCAmelCase=None							,			**lowerCAmelCase							,			):





      """simple docstring"""


      if map_location not in [None, "cpu", "on_device"]:
            raise TypeError(
                """Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"""						)
      if map_location is None:
            _lowerCAmelCase							  =							"""cpu"""
      elif map_location == "on_device":
            _lowerCAmelCase							  =							PartialState().device
      # Model states
      for i, model in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            models[i].load_state_dict(torch.load(lowerCAmelCase							,			map_location=lowerCAmelCase						)							,			**lowerCAmelCase						)
      logger.info("""All model weights loaded successfully"""						)

      # Optimizer states
      for i, opt in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            _lowerCAmelCase							  =							torch.load(lowerCAmelCase							,			map_location=lowerCAmelCase						)
            optimizers[i].load_state_dict(lowerCAmelCase						)
      logger.info("""All optimizer states loaded successfully"""						)

      # Scheduler states
      for i, scheduler in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            scheduler.load_state_dict(torch.load(lowerCAmelCase						)						)
      logger.info("""All scheduler states loaded successfully"""						)

      # GradScaler state
      if scaler is not None:
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            scaler.load_state_dict(torch.load(lowerCAmelCase						)						)
            logger.info("""GradScaler state loaded successfully"""						)

      # Random states
      try:
            _lowerCAmelCase							  =							torch.load(os.path.join(lowerCAmelCase							,			f"{RNG_STATE_NAME}_{process_index}.pkl"						)						)
            random.setstate(states["""random_state"""]						)
            np.random.set_state(states["""numpy_random_seed"""]						)
            torch.set_rng_state(states["""torch_manual_seed"""]						)
            if is_xpu_available():
                  torch.xpu.set_rng_state_all(states["""torch_xpu_manual_seed"""]						)
            else:
                  torch.cuda.set_rng_state_all(states["""torch_cuda_manual_seed"""]						)
            if is_tpu_available():
                  xm.set_rng_state(states["""xm_seed"""]						)
            logger.info("""All random states loaded successfully"""						)
      except Exception:
            logger.info("""Could not load random states"""						)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = 0						):





      """simple docstring"""


      _lowerCAmelCase							  =							Path(lowerCAmelCase						) / f"custom_checkpoint_{index}.pkl"
      logger.info(f"Saving the state of {get_pretty_name(lowerCAmelCase						)} to {save_location}"						)
      torch.save(obj.state_dict()							,			lowerCAmelCase						)





def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = 0						):





      """simple docstring"""


      _lowerCAmelCase							  =							f"{path}/custom_checkpoint_{index}.pkl"
      logger.info(f"Loading the state of {get_pretty_name(lowerCAmelCase						)} from {load_location}"						)
      obj.load_state_dict(torch.load(lowerCAmelCase							,			map_location="""cpu"""						)						)






