





'''simple docstring'''






import gc
import random
import unittest

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

from diffusers import AutoencoderKL, CycleDiffusionPipeline, DDIMScheduler, UNetaDConditionModel
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, skip_mps

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 			UpperCAmelCase					(      snake_case_							,		snake_case_							,		unittest.TestCase       ):
      _lowercase:      List[str]											=       CycleDiffusionPipeline
      _lowercase:      str											=       TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
          '''negative_prompt''',
          '''height''',
          '''width''',
          '''negative_prompt_embeds''',
      }
      _lowercase:      Union[str, Any]											=       PipelineTesterMixin.required_optional_params - {'''latents'''}
      _lowercase:      List[Any]											=       TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS.union({'''source_prompt'''}       )
      _lowercase:      Optional[Any]											=       IMAGE_TO_IMAGE_IMAGE_PARAMS
      _lowercase:      Optional[Any]											=       IMAGE_TO_IMAGE_IMAGE_PARAMS




      def     lowercase__   (      self      :			Tuple  )    ->				Union[str, Any]:
            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      ,      )
            _lowerCAmelCase												=							DDIMScheduler(
                beta_start=0.0_00_85      ,      beta_end=0.0_12      ,      beta_schedule="""scaled_linear"""      ,      num_train_timesteps=10_00      ,      clip_sample=_lowerCAmelCase      ,      set_alpha_to_one=_lowerCAmelCase      ,      )
            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      ,      )
            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      ,      )
            _lowerCAmelCase												=							CLIPTextModel(_lowerCAmelCase  )
            _lowerCAmelCase												=							CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

            _lowerCAmelCase												=							{
                """unet""": unet,
                """scheduler""": scheduler,
                """vae""": vae,
                """text_encoder""": text_encoder,
                """tokenizer""": tokenizer,
                """safety_checker""": None,
                """feature_extractor""": None,
            }
            return components




      def     lowercase__   (      self      :			List[Any]      ,      __snake_case      :			Tuple      ,      __snake_case      :			Union[str, Any]=0  )    ->				int:
            _lowerCAmelCase												=							floats_tensor((1, 3, 32, 32)      ,      rng=random.Random(_lowerCAmelCase  )  ).to(_lowerCAmelCase  )
            _lowerCAmelCase												=							image / 2 + 0.5
            if str(_lowerCAmelCase  ).startswith("""mps"""  ):
                  _lowerCAmelCase												=							torch.manual_seed(_lowerCAmelCase  )
            else:
                  _lowerCAmelCase												=							torch.Generator(device=_lowerCAmelCase  ).manual_seed(_lowerCAmelCase  )
            _lowerCAmelCase												=							{
                """prompt""": """An astronaut riding an elephant""",
                """source_prompt""": """An astronaut riding a horse""",
                """image""": image,
                """generator""": generator,
                """num_inference_steps""": 2,
                """eta""": 0.1,
                """strength""": 0.8,
                """guidance_scale""": 3,
                """source_guidance_scale""": 1,
                """output_type""": """numpy""",
            }
            return inputs




      def     lowercase__   (      self      :			Any  )    ->				str:
            _lowerCAmelCase												=							"""cpu"""  # ensure determinism for the device-dependent torch.Generator

            _lowerCAmelCase												=							self.get_dummy_components()
            _lowerCAmelCase												=							CycleDiffusionPipeline(**_lowerCAmelCase  )
            _lowerCAmelCase												=							pipe.to(_lowerCAmelCase  )
            pipe.set_progress_bar_config(disable=_lowerCAmelCase  )

            _lowerCAmelCase												=							self.get_dummy_inputs(_lowerCAmelCase  )
            _lowerCAmelCase												=							pipe(**_lowerCAmelCase  )
            _lowerCAmelCase												=							output.images

            _lowerCAmelCase												=							images[0, -3:, -3:, -1]

            assert images.shape == (1, 32, 32, 3)
            _lowerCAmelCase												=							np.array([0.44_59, 0.49_43, 0.45_44, 0.66_43, 0.54_74, 0.43_27, 0.57_01, 0.59_59, 0.51_79]  )

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




      @unittest.skipIf(torch_device != """cuda"""      ,      """This test requires a GPU"""  )
      def     lowercase__   (      self      :			List[str]  )    ->				Tuple:
            _lowerCAmelCase												=							self.get_dummy_components()
            for name, module in components.items():
                  if hasattr(_lowerCAmelCase      ,      """half"""  ):
                        _lowerCAmelCase												=							module.half()
            _lowerCAmelCase												=							CycleDiffusionPipeline(**_lowerCAmelCase  )
            _lowerCAmelCase												=							pipe.to(_lowerCAmelCase  )
            pipe.set_progress_bar_config(disable=_lowerCAmelCase  )

            _lowerCAmelCase												=							self.get_dummy_inputs(_lowerCAmelCase  )
            _lowerCAmelCase												=							pipe(**_lowerCAmelCase  )
            _lowerCAmelCase												=							output.images

            _lowerCAmelCase												=							images[0, -3:, -3:, -1]

            assert images.shape == (1, 32, 32, 3)
            _lowerCAmelCase												=							np.array([0.35_06, 0.45_43, 0.4_46, 0.45_75, 0.51_95, 0.41_55, 0.52_73, 0.5_18, 0.41_16]  )

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




      @skip_mps
      def     lowercase__   (      self      :			Optional[int]  )    ->				int:
            return super().test_save_load_local()




      @unittest.skip("""non-deterministic pipeline"""  )
      def     lowercase__   (      self      :			List[str]  )    ->				Optional[int]:
            return super().test_inference_batch_single_identical()




      @skip_mps
      def     lowercase__   (      self      :			int  )    ->				Any:
            return super().test_dict_tuple_outputs_equivalent()




      @skip_mps
      def     lowercase__   (      self      :			Optional[Any]  )    ->				Optional[int]:
            return super().test_save_load_optional_components()




      @skip_mps
      def     lowercase__   (      self      :			Any  )    ->				List[str]:
            return super().test_attention_slicing_forward_pass()




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




      def     lowercase__   (      self      :			Optional[Any]  )    ->				Optional[int]:
            super().tearDown()
            gc.collect()
            torch.cuda.empty_cache()




      def     lowercase__   (      self      :			Optional[Any]  )    ->				List[str]:
            _lowerCAmelCase												=							load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/cycle-diffusion/black_colored_car.png"""  )
            _lowerCAmelCase												=							load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car_fp16.npy"""  )
            _lowerCAmelCase												=							init_image.resize((5_12, 5_12)  )

            _lowerCAmelCase												=							"""CompVis/stable-diffusion-v1-4"""
            _lowerCAmelCase												=							DDIMScheduler.from_pretrained(_lowerCAmelCase      ,      subfolder="""scheduler"""  )
            _lowerCAmelCase												=							CycleDiffusionPipeline.from_pretrained(
                _lowerCAmelCase      ,      scheduler=_lowerCAmelCase      ,      safety_checker=_lowerCAmelCase      ,      torch_dtype=torch.floataa      ,      revision="""fp16"""  )

            pipe.to(_lowerCAmelCase  )
            pipe.set_progress_bar_config(disable=_lowerCAmelCase  )
            pipe.enable_attention_slicing()

            _lowerCAmelCase												=							"""A black colored car"""
            _lowerCAmelCase												=							"""A blue colored car"""

            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							pipe(
                prompt=_lowerCAmelCase      ,      source_prompt=_lowerCAmelCase      ,      image=_lowerCAmelCase      ,      num_inference_steps=1_00      ,      eta=0.1      ,      strength=0.85      ,      guidance_scale=3      ,      source_guidance_scale=1      ,      generator=_lowerCAmelCase      ,      output_type="""np"""      ,      )
            _lowerCAmelCase												=							output.images

            # the values aren't exactly equal, but the images look the same visually
            assert np.abs(image - expected_image  ).max() < 5E-1




      def     lowercase__   (      self      :			Any  )    ->				Tuple:
            _lowerCAmelCase												=							load_image(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
                """/cycle-diffusion/black_colored_car.png"""  )
            _lowerCAmelCase												=							load_numpy(
                """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/cycle-diffusion/blue_colored_car.npy"""  )
            _lowerCAmelCase												=							init_image.resize((5_12, 5_12)  )

            _lowerCAmelCase												=							"""CompVis/stable-diffusion-v1-4"""
            _lowerCAmelCase												=							DDIMScheduler.from_pretrained(_lowerCAmelCase      ,      subfolder="""scheduler"""  )
            _lowerCAmelCase												=							CycleDiffusionPipeline.from_pretrained(_lowerCAmelCase      ,      scheduler=_lowerCAmelCase      ,      safety_checker=_lowerCAmelCase  )

            pipe.to(_lowerCAmelCase  )
            pipe.set_progress_bar_config(disable=_lowerCAmelCase  )
            pipe.enable_attention_slicing()

            _lowerCAmelCase												=							"""A black colored car"""
            _lowerCAmelCase												=							"""A blue colored car"""

            _lowerCAmelCase												=							torch.manual_seed(0  )
            _lowerCAmelCase												=							pipe(
                prompt=_lowerCAmelCase      ,      source_prompt=_lowerCAmelCase      ,      image=_lowerCAmelCase      ,      num_inference_steps=1_00      ,      eta=0.1      ,      strength=0.85      ,      guidance_scale=3      ,      source_guidance_scale=1      ,      generator=_lowerCAmelCase      ,      output_type="""np"""      ,      )
            _lowerCAmelCase												=							output.images

            assert np.abs(image - expected_image  ).max() < 2E-2













'''simple docstring'''






import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
A__    :       Any              =logging.getLogger()

A__    :       int              =logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)




class 			UpperCAmelCase					(      snake_case_       ):




      def     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			Any  )    ->				int:
            os.makedirs(__snake_case      ,      exist_ok=__snake_case  )
            _lowerCAmelCase												=							{"""source""": """What is love ?""", """target""": """life"""}
            _lowerCAmelCase												=							{"""train""": 12, """val""": 2, """test""": 2}
            for split in ["train", "test", "val"]:
                  for field in ["source", "target"]:
                        _lowerCAmelCase												=							"""\n""".join([contents[field]] * n_lines[split]  )
                        with open(os.path.join(__snake_case      ,      f"{split}.{field}"  )      ,      """w"""  ) as f:
                              f.write(__snake_case  )




      def     lowercase__   (      self      :			Tuple      ,      __snake_case      :			int      ,      __snake_case      :			str = "pytorch"  )    ->				int:
            _lowerCAmelCase												=							self.get_auto_remove_tmp_dir()
            _lowerCAmelCase												=							os.path.join(__snake_case      ,      """output"""  )
            _lowerCAmelCase												=							os.path.join(__snake_case      ,      """data"""  )
            self._create_dummy_data(data_dir=__snake_case  )

            _lowerCAmelCase												=							f"\n                --data_dir {data_dir} \\n                --output_dir {output_dir} \\n                --model_name_or_path facebook/rag-sequence-base \\n                --model_type rag_sequence \\n                --do_train \\n                --do_predict \\n                --n_val -1 \\n                --val_check_interval 1.0 \\n                --train_batch_size 2 \\n                --eval_batch_size 1 \\n                --max_source_length 25 \\n                --max_target_length 25 \\n                --val_max_target_length 25 \\n                --test_max_target_length 25 \\n                --label_smoothing 0.1 \\n                --dropout 0.1 \\n                --attention_dropout 0.1 \\n                --weight_decay 0.001 \\n                --adam_epsilon 1e-08 \\n                --max_grad_norm 0.1 \\n                --lr_scheduler polynomial \\n                --learning_rate 3e-04 \\n                --num_train_epochs 1 \\n                --warmup_steps 4 \\n                --gradient_accumulation_steps 1 \\n                --distributed-port 8787 \\n                --use_dummy_dataset 1 \\n                --distributed_retriever {distributed_retriever} \\n            ".split()

            if gpus > 0:
                  testargs.append(f"--gpus={gpus}"  )
                  if is_apex_available():
                        testargs.append("""--fp16"""  )
            else:
                  testargs.append("""--gpus=0"""  )
                  testargs.append("""--distributed_backend=ddp_cpu"""  )
                  testargs.append("""--num_processes=2"""  )

            _lowerCAmelCase												=							[sys.executable, str(Path(finetune_rag.__file__  ).resolve()  )] + testargs
            execute_subprocess_async(__snake_case      ,      env=self.get_env()  )

            _lowerCAmelCase												=							os.path.join(__snake_case      ,      """metrics.json"""  )
            with open(__snake_case  ) as f:
                  _lowerCAmelCase												=							json.load(__snake_case  )
            return result




      @require_torch_gpu
      def     lowercase__   (      self      :			Dict  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							self._run_finetune(gpus=1  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )




      @require_torch_multi_gpu
      def     lowercase__   (      self      :			int  )    ->				Dict:
            _lowerCAmelCase												=							self._run_finetune(gpus=2  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )




      @require_torch_gpu
      @require_ray
      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							self._run_finetune(gpus=1      ,      distributed_retriever="""ray"""  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )




      @require_torch_multi_gpu
      @require_ray
      def     lowercase__   (      self      :			int  )    ->				List[str]:
            _lowerCAmelCase												=							self._run_finetune(gpus=1      ,      distributed_retriever="""ray"""  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )






