

"""simple docstring"""





import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image

from diffusers import (
    DDIMScheduler,
    KandinskyVaaImgaImgPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 					_lowerCAmelCase  (	lowercase     ,unittest.TestCase   ):







 """simple docstring"""







 __UpperCAmelCase	:					int   		=	KandinskyVaaImgaImgPipeline
 __UpperCAmelCase	:					Optional[Any]   		=	["image_embeds", "negative_image_embeds", "image"]
 __UpperCAmelCase	:					int   		=	[
     "image_embeds",
     "negative_image_embeds",
     "image",
 ]
 __UpperCAmelCase	:					List[str]   		=	[
     "generator",
     "height",
     "width",
     "strength",
     "guidance_scale",
     "num_inference_steps",
     "return_dict",
     "guidance_scale",
     "num_images_per_prompt",
     "output_type",
     "return_dict",
 ]
 __UpperCAmelCase	:					List[Any]   		=	False


 @property
 def 					_lowercase (			self							:			Optional[int]	):
    return 3_2


 @property
 def 					_lowercase (			self							:			Union[str, Any]	):
    return 3_2


 @property
 def 					_lowercase (			self							:			Optional[int]	):
    return self.time_input_dim


 @property
 def 					_lowercase (			self							:			Any	):
    return self.time_input_dim * 4


 @property
 def 					_lowercase (			self							:			Any	):
    return 1_0_0


 @property
 def 					_lowercase (			self							:			Any	):
    torch.manual_seed(0	)

    __lowercase								=						{
        "in_channels": 4,
        # Out channels is double in channels because predicts mean and variance
        "out_channels": 8,
        "addition_embed_type": "image",
        "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
        "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
        "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
        "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
        "layers_per_block": 1,
        "encoder_hid_dim": self.text_embedder_hidden_size,
        "encoder_hid_dim_type": "image_proj",
        "cross_attention_dim": self.cross_attention_dim,
        "attention_head_dim": 4,
        "resnet_time_scale_shift": "scale_shift",
        "class_embed_type": None,
    }

    __lowercase								=						UNetaDConditionModel(**UpperCAmelCase__	)
    return model


 @property
 def 					_lowercase (			self							:			Optional[Any]	):
    return {
        "block_out_channels": [3_2, 6_4],
        "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
        "in_channels": 3,
        "latent_channels": 4,
        "layers_per_block": 1,
        "norm_num_groups": 8,
        "norm_type": "spatial",
        "num_vq_embeddings": 1_2,
        "out_channels": 3,
        "up_block_types": [
            "AttnUpDecoderBlock2D",
            "UpDecoderBlock2D",
        ],
        "vq_embed_dim": 4,
    }


 @property
 def 					_lowercase (			self							:			Optional[Any]	):
    torch.manual_seed(0	)
    __lowercase								=						VQModel(**self.dummy_movq_kwargs	)
    return model


 def 					_lowercase (			self							:			Union[str, Any]	):
    __lowercase								=						self.dummy_unet
    __lowercase								=						self.dummy_movq

    __lowercase								=						{
        "num_train_timesteps": 1_0_0_0,
        "beta_schedule": "linear",
        "beta_start": 0.00_085,
        "beta_end": 0.012,
        "clip_sample": False,
        "set_alpha_to_one": False,
        "steps_offset": 0,
        "prediction_type": "epsilon",
        "thresholding": False,
    }

    __lowercase								=						DDIMScheduler(**UpperCAmelCase__	)

    __lowercase								=						{
        "unet": unet,
        "scheduler": scheduler,
        "movq": movq,
    }

    return components


 def 					_lowercase (			self							:			Union[str, Any],					UpperCAmelCase__							:			List[Any],					UpperCAmelCase__							:			List[Any]=0	):
    __lowercase								=						floats_tensor((1, self.text_embedder_hidden_size),					rng=random.Random(UpperCAmelCase__	)	).to(UpperCAmelCase__	)
    __lowercase								=						floats_tensor((1, self.text_embedder_hidden_size),					rng=random.Random(seed + 1	)	).to(
        UpperCAmelCase__	)
    # create init_image
    __lowercase								=						floats_tensor((1, 3, 6_4, 6_4),					rng=random.Random(UpperCAmelCase__	)	).to(UpperCAmelCase__	)
    __lowercase								=						image.cpu().permute(0,					2,					3,					1	)[0]
    __lowercase								=						Image.fromarray(np.uinta(UpperCAmelCase__	)	).convert("RGB"	).resize((2_5_6, 2_5_6)	)

    if str(UpperCAmelCase__	).startswith("mps"	):
       __lowercase								=						torch.manual_seed(UpperCAmelCase__	)
    else:
       __lowercase								=						torch.Generator(device=UpperCAmelCase__	).manual_seed(UpperCAmelCase__	)
    __lowercase								=						{
        "image": init_image,
        "image_embeds": image_embeds,
        "negative_image_embeds": negative_image_embeds,
        "generator": generator,
        "height": 6_4,
        "width": 6_4,
        "num_inference_steps": 1_0,
        "guidance_scale": 7.0,
        "strength": 0.2,
        "output_type": "np",
    }
    return inputs



 def 					_lowercase (			self							:			Optional[Any]	):
    __lowercase								=						"cpu"

    __lowercase								=						self.get_dummy_components()

    __lowercase								=						self.pipeline_class(**UpperCAmelCase__	)
    __lowercase								=						pipe.to(UpperCAmelCase__	)

    pipe.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						pipe(**self.get_dummy_inputs(UpperCAmelCase__	)	)
    __lowercase								=						output.images

    __lowercase								=						pipe(
        **self.get_dummy_inputs(UpperCAmelCase__	),					return_dict=UpperCAmelCase__,					)[0]

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

    assert image.shape == (1, 6_4, 6_4, 3)

    __lowercase								=						np.array(
        [0.6_199_778, 0.63_984_406, 0.46_145_785, 0.62_944_984, 0.5_622_215, 0.47_306_132, 0.47_441_456, 0.4_607_606, 0.48_719_263]	)
    assert (
        np.abs(image_slice.flatten() - expected_slice	).max() < 1E-2
    ), F""" expected_slice {expected_slice}, but got {image_slice.flatten()}"""
    assert (
        np.abs(image_from_tuple_slice.flatten() - expected_slice	).max() < 1E-2
    ), F""" expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}"""







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







 """simple docstring"""







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



 def 					_lowercase (			self							:			Dict	):
    __lowercase								=						load_numpy(
        "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
        "/kandinskyv22/kandinskyv22_img2img_frog.npy"	)

    __lowercase								=						load_image(
        "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main" "/kandinsky/cat.png"	)
    __lowercase								=						"A red cartoon frog, 4k"

    __lowercase								=						KandinskyVaaPriorPipeline.from_pretrained(
        "kandinsky-community/kandinsky-2-2-prior",					torch_dtype=torch.floataa	)
    pipe_prior.to(UpperCAmelCase__	)

    __lowercase								=						KandinskyVaaImgaImgPipeline.from_pretrained(
        "kandinsky-community/kandinsky-2-2-decoder",					torch_dtype=torch.floataa	)
    __lowercase								=						pipeline.to(UpperCAmelCase__	)

    pipeline.set_progress_bar_config(disable=UpperCAmelCase__	)

    __lowercase								=						torch.Generator(device="cpu"	).manual_seed(0	)
    __lowercase		,__lowercase								=						pipe_prior(
        UpperCAmelCase__,					generator=UpperCAmelCase__,					num_inference_steps=5,					negative_prompt="",					).to_tuple()

    __lowercase								=						pipeline(
        image=UpperCAmelCase__,					image_embeds=UpperCAmelCase__,					negative_image_embeds=UpperCAmelCase__,					generator=UpperCAmelCase__,					num_inference_steps=1_0_0,					height=7_6_8,					width=7_6_8,					strength=0.2,					output_type="np",					)

    __lowercase								=						output.images[0]

    assert image.shape == (7_6_8, 7_6_8, 3)

    assert_mean_pixel_difference(UpperCAmelCase__,					UpperCAmelCase__	)







"""simple docstring"""





import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
   from sagemaker import Session, TrainingJobAnalytics
   from sagemaker.huggingface import HuggingFace




@pytest.mark.skipif(
    literal_eval(os.getenv("TEST_SAGEMAKER"     ,"False"   )   ) is not True     ,reason="Skipping test because should only be run when releasing minor transformers version"     ,)
@pytest.mark.usefixtures("sm_env"   )
@parameterized_class(
    [
        {
            "framework": "pytorch",
            "script": "run_glue_model_parallelism.py",
            "model_name_or_path": "roberta-large",
            "instance_type": "ml.p3dn.24xlarge",
            "results": {"train_runtime": 1_6_0_0, "eval_accuracy": 0.3, "eval_loss": 1.2},
        },
        {
            "framework": "pytorch",
            "script": "run_glue.py",
            "model_name_or_path": "roberta-large",
            "instance_type": "ml.p3dn.24xlarge",
            "results": {"train_runtime": 1_6_0_0, "eval_accuracy": 0.3, "eval_loss": 1.2},
        },
    ]   )
class 					_lowerCAmelCase  (	unittest.TestCase   ):







 """simple docstring"""







 def 					_lowercase (			self							:			Optional[int]	):
    if self.framework == "pytorch":
       subprocess.run(
           F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split(),					encoding="utf-8",					check=UpperCAmelCase__,					)
    assert hasattr(self,					"env"	)


 def 					_lowercase (			self							:			str,					UpperCAmelCase__							:			List[Any]	):
    # configuration for running training on smdistributed Model Parallel
    __lowercase								=						{
        "enabled": True,
        "processes_per_host": 8,
    }
    __lowercase								=						{
        "enabled": True,
        "parameters": {
            "microbatches": 4,
            "placement_strategy": "spread",
            "pipeline": "interleaved",
            "optimize": "speed",
            "partitions": 4,
            "ddp": True,
        },
    }

    __lowercase								=						{"smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options}

    __lowercase								=						"trainer" if self.script == "run_glue.py" else "smtrainer"
    # creates estimator
    return HuggingFace(
        entry_point=self.script,					source_dir=self.env.test_path,					role=self.env.role,					image_uri=self.env.image_uri,					base_job_name=F"""{self.env.base_job_name}-{instance_count}-smp-{name_extension}""",					instance_count=UpperCAmelCase__,					instance_type=self.instance_type,					debugger_hook_config=UpperCAmelCase__,					hyperparameters={
            **self.env.hyperparameters,
            "model_name_or_path": self.model_name_or_path,
            "max_steps": 5_0_0,
        },					metric_definitions=self.env.metric_definitions,					distribution=UpperCAmelCase__,					py_version="py36",					)


 def 					_lowercase (			self							:			Tuple,					UpperCAmelCase__							:			int	):
    TrainingJobAnalytics(UpperCAmelCase__	).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""	)



 @parameterized.expand([(1,)]	)
 def 					_lowercase (			self							:			str,					UpperCAmelCase__							:			Union[str, Any]	):
    # create estimator
    __lowercase								=						self.create_estimator(UpperCAmelCase__	)

    # run training
    estimator.fit()

    # result dataframe
    __lowercase								=						TrainingJobAnalytics(estimator.latest_training_job.name	).dataframe()

    # extract kpis
    __lowercase								=						list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]	)
    __lowercase								=						list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]	)
    # get train time from SageMaker job, this includes starting, preprocessing, stopping
    __lowercase								=						(
        Session().describe_training_job(estimator.latest_training_job.name	).get("TrainingTimeInSeconds",					9_9_9_9_9_9	)
    )

    # assert kpis
    assert train_runtime <= self.results["train_runtime"]
    assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy	)
    assert all(t <= self.results["eval_loss"] for t in eval_loss	)

    # dump tests result into json file to share in PR
    with open(F"""{estimator.latest_training_job.name}.json""",					"w"	) as outfile:
       json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss},					UpperCAmelCase__	)




