


import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







class lowerCAmelCase_						( _A					,						_A					,						_A					,						unittest.TestCase       ):

 '''simple docstring'''

 __snake_case												=  StableDiffusionImageVariationPipeline
 __snake_case												=  IMAGE_VARIATION_PARAMS
 __snake_case												=  IMAGE_VARIATION_BATCH_PARAMS
 __snake_case												=  frozenset([]       )
 # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 __snake_case												=  frozenset([]       )



 def 						UpperCamelCase__     (	self					):
       torch.manual_seed(0					)
       snake_case_       =      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		,      )
       snake_case_       =      PNDMScheduler(skip_prk_steps=__SCREAMING_SNAKE_CASE					)
       torch.manual_seed(0					)
       snake_case_       =      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					)
       snake_case_       =      CLIPVisionConfig(
           hidden_size=32		,      projection_dim=32		,      intermediate_size=37		,      layer_norm_eps=1E-05		,      num_attention_heads=4		,      num_hidden_layers=5		,      image_size=32		,      patch_size=4		,      )
       snake_case_       =      CLIPVisionModelWithProjection(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      CLIPImageProcessor(crop_size=32		,      size=32					)

       snake_case_       =      {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''image_encoder''': image_encoder,
           '''feature_extractor''': feature_extractor,
           '''safety_checker''': None,
       }
       return components



 def 						UpperCamelCase__     (	self		,      _UpperCAmelCase		,      _UpperCAmelCase=0					):
       snake_case_       =      floats_tensor((1, 3, 32, 32)		,      rng=random.Random(__SCREAMING_SNAKE_CASE					)					)
       snake_case_       =      image.cpu().permute(0		,      2		,      3		,      1					)[0]
       snake_case_       =      Image.fromarray(np.uinta(__SCREAMING_SNAKE_CASE					)					).convert('''RGB'''					).resize((32, 32)					)
       if str(__SCREAMING_SNAKE_CASE					).startswith('''mps'''					):
             snake_case_       =      torch.manual_seed(__SCREAMING_SNAKE_CASE					)
       else:
             snake_case_       =      torch.Generator(device=__SCREAMING_SNAKE_CASE					).manual_seed(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      {
           '''image''': image,
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }
       return inputs



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       snake_case_       =      self.get_dummy_components()
       snake_case_       =      StableDiffusionImageVariationPipeline(**__SCREAMING_SNAKE_CASE					)
       snake_case_       =      sd_pipe.to(__SCREAMING_SNAKE_CASE					)
       sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)

       snake_case_       =      self.get_dummy_inputs(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      sd_pipe(**__SCREAMING_SNAKE_CASE					).images
       snake_case_       =      image[0, -3:, -3:, -1]

       assert image.shape == (1, 64, 64, 3)
       snake_case_       =      np.array([0.5_239, 0.5_723, 0.4_796, 0.5_049, 0.5_550, 0.4_685, 0.5_329, 0.4_891, 0.4_921]					)

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



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      '''cpu'''  # ensure determinism for the device-dependent torch.Generator
       snake_case_       =      self.get_dummy_components()
       snake_case_       =      StableDiffusionImageVariationPipeline(**__SCREAMING_SNAKE_CASE					)
       snake_case_       =      sd_pipe.to(__SCREAMING_SNAKE_CASE					)
       sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)

       snake_case_       =      self.get_dummy_inputs(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      2 * [inputs['''image''']]
       snake_case_       =      sd_pipe(**__SCREAMING_SNAKE_CASE					)

       snake_case_       =      output.images

       snake_case_       =      image[-1, -3:, -3:, -1]

       assert image.shape == (2, 64, 64, 3)
       snake_case_       =      np.array([0.6_892, 0.5_637, 0.5_836, 0.5_771, 0.6_254, 0.6_409, 0.5_580, 0.5_569, 0.5_289]					)

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





 def 						UpperCamelCase__     (	self					):
       super().test_inference_batch_single_identical(expected_max_diff=3E-3					)







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

 '''simple docstring'''

 def 						UpperCamelCase__     (	self					):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 def 						UpperCamelCase__     (	self		,      _UpperCAmelCase		,      _UpperCAmelCase="cpu"		,      _UpperCAmelCase=torch.floataa		,      _UpperCAmelCase=0					):
       snake_case_       =      torch.Generator(device=__SCREAMING_SNAKE_CASE					).manual_seed(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      load_image(
           '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
           '''/stable_diffusion_imgvar/input_image_vermeer.png'''					)
       snake_case_       =      np.random.RandomState(__SCREAMING_SNAKE_CASE					).standard_normal((1, 4, 64, 64)					)
       snake_case_       =      torch.from_numpy(__SCREAMING_SNAKE_CASE					).to(device=__SCREAMING_SNAKE_CASE		,      dtype=__SCREAMING_SNAKE_CASE					)
       snake_case_       =      {
           '''image''': init_image,
           '''latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 3,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      StableDiffusionImageVariationPipeline.from_pretrained(
           '''lambdalabs/sd-image-variations-diffusers'''		,      safety_checker=__SCREAMING_SNAKE_CASE					)
       snake_case_       =      sd_pipe.to(__SCREAMING_SNAKE_CASE					)
       sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)

       snake_case_       =      self.get_inputs(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      sd_pipe(**__SCREAMING_SNAKE_CASE					).images
       snake_case_       =      image[0, -3:, -3:, -1].flatten()

       assert image.shape == (1, 5_12, 5_12, 3)
       snake_case_       =      np.array([0.84_491, 0.90_789, 0.75_708, 0.78_734, 0.83_485, 0.70_099, 0.66_938, 0.68_727, 0.61_379]					)
       assert np.abs(image_slice - expected_slice					).max() < 6E-3



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      0

       def callback_fn(_UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase					) -> None:
             snake_case_       =      True
             nonlocal number_of_steps
             number_of_steps += 1
             if step == 1:
                   snake_case_       =      latents.detach().cpu().numpy()
                   assert latents.shape == (1, 4, 64, 64)
                   snake_case_       =      latents[0, -3:, -3:, -1]
                   snake_case_       =      np.array(
                       [-0.1_621, 0.2_837, -0.7_979, -0.1_221, -1.3_057, 0.7_681, -2.1_191, 0.0_464, 1.6_309]					)

                   assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2
             elif step == 2:
                   snake_case_       =      latents.detach().cpu().numpy()
                   assert latents.shape == (1, 4, 64, 64)
                   snake_case_       =      latents[0, -3:, -3:, -1]
                   snake_case_       =      np.array([0.6_299, 1.7_500, 1.1_992, -2.1_582, -1.8_994, 0.7_334, -0.7_090, 1.0_137, 1.5_273]					)

                   assert np.abs(latents_slice.flatten() - expected_slice					).max() < 5E-2

       snake_case_       =      False

       snake_case_       =      StableDiffusionImageVariationPipeline.from_pretrained(
           '''fusing/sd-image-variations-diffusers'''		,      safety_checker=__SCREAMING_SNAKE_CASE		,      torch_dtype=torch.floataa		,      )
       pipe.to(__SCREAMING_SNAKE_CASE					)
       pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)
       pipe.enable_attention_slicing()

       snake_case_       =      self.get_inputs(__SCREAMING_SNAKE_CASE		,      dtype=torch.floataa					)
       pipe(**__SCREAMING_SNAKE_CASE		,      callback=__SCREAMING_SNAKE_CASE		,      callback_steps=1					)
       assert callback_fn.has_been_called
       assert number_of_steps == inputs["num_inference_steps"]





 def 						UpperCamelCase__     (	self					):
       torch.cuda.empty_cache()
       torch.cuda.reset_max_memory_allocated()
       torch.cuda.reset_peak_memory_stats()

       snake_case_       =      '''fusing/sd-image-variations-diffusers'''
       snake_case_       =      StableDiffusionImageVariationPipeline.from_pretrained(
           __SCREAMING_SNAKE_CASE		,      safety_checker=__SCREAMING_SNAKE_CASE		,      torch_dtype=torch.floataa					)
       snake_case_       =      pipe.to(__SCREAMING_SNAKE_CASE					)
       pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)
       pipe.enable_attention_slicing(1					)
       pipe.enable_sequential_cpu_offload()

       snake_case_       =      self.get_inputs(__SCREAMING_SNAKE_CASE		,      dtype=torch.floataa					)
       snake_case_       =      pipe(**__SCREAMING_SNAKE_CASE					)

       snake_case_       =      torch.cuda.max_memory_allocated()
       # make sure that less than 2.6 GB is allocated
       assert mem_bytes < 2.6 * 10**9







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

 '''simple docstring'''

 def 						UpperCamelCase__     (	self					):
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()



 def 						UpperCamelCase__     (	self		,      _UpperCAmelCase		,      _UpperCAmelCase="cpu"		,      _UpperCAmelCase=torch.floataa		,      _UpperCAmelCase=0					):
       snake_case_       =      torch.Generator(device=__SCREAMING_SNAKE_CASE					).manual_seed(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      load_image(
           '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
           '''/stable_diffusion_imgvar/input_image_vermeer.png'''					)
       snake_case_       =      np.random.RandomState(__SCREAMING_SNAKE_CASE					).standard_normal((1, 4, 64, 64)					)
       snake_case_       =      torch.from_numpy(__SCREAMING_SNAKE_CASE					).to(device=__SCREAMING_SNAKE_CASE		,      dtype=__SCREAMING_SNAKE_CASE					)
       snake_case_       =      {
           '''image''': init_image,
           '''latents''': latents,
           '''generator''': generator,
           '''num_inference_steps''': 50,
           '''guidance_scale''': 7.5,
           '''output_type''': '''numpy''',
       }
       return inputs



 def 						UpperCamelCase__     (	self					):
       snake_case_       =      StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers'''					)
       sd_pipe.to(__SCREAMING_SNAKE_CASE					)
       sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)

       snake_case_       =      self.get_inputs(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      sd_pipe(**__SCREAMING_SNAKE_CASE					).images[0]

       snake_case_       =      load_numpy(
           '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
           '''/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy'''					)
       snake_case_       =      np.abs(expected_image - image					).max()
       assert max_diff < 1E-3





 def 						UpperCamelCase__     (	self					):
       snake_case_       =      StableDiffusionImageVariationPipeline.from_pretrained('''fusing/sd-image-variations-diffusers'''					)
       snake_case_       =      DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config					)
       sd_pipe.to(__SCREAMING_SNAKE_CASE					)
       sd_pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE					)

       snake_case_       =      self.get_inputs(__SCREAMING_SNAKE_CASE					)
       snake_case_       =      25
       snake_case_       =      sd_pipe(**__SCREAMING_SNAKE_CASE					).images[0]

       snake_case_       =      load_numpy(
           '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
           '''/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy'''					)
       snake_case_       =      np.abs(expected_image - image					).max()
       assert max_diff < 1E-3


"""simple docstring"""




import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Image
from datasets.packaged_modules.text.text import Text

from ..utils import require_pil
@pytest.fixture
def    __SCREAMING_SNAKE_CASE						(						lowercase__      ):
					"""simple docstring"""
					A			=					tmp_path / "text.txt"
					A			=					textwrap.dedent(
					    "\\n        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.\n        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.\n        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.\n        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.\n\n        Second paragraph:\n        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.\n        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.\n        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.\n        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.\n        "      )
					with open(lowercase__			,							"w"			,							encoding="utf-8"      ) as f:
										f.write(lowercase__      )
					return str(lowercase__      )
@pytest.fixture
def    __SCREAMING_SNAKE_CASE						(						lowercase__			,							lowercase__      ):
					"""simple docstring"""
					A			=					tmp_path / "text_with_image.txt"
					with open(lowercase__			,							"w"			,							encoding="utf-8"      ) as f:
										f.write(lowercase__      )
					return str(lowercase__      )
@pytest.mark.parametrize("keep_linebreaks"			,							[True, False]      )
def    __SCREAMING_SNAKE_CASE						(						lowercase__			,							lowercase__      ):
					"""simple docstring"""
					with open(lowercase__			,							encoding="utf-8"      ) as f:
										A			=					f.read().splitlines(keepends=lowercase__      )
					A			=					Text(keep_linebreaks=lowercase__			,							encoding="utf-8"      )
					A			=					text._generate_tables([[text_file]]      )
					A			=					pa.concat_tables([table for _, table in generator]      ).to_pydict()["text"]
					assert generated_content == expected_content
@require_pil
def    __SCREAMING_SNAKE_CASE						(						lowercase__      ):
					"""simple docstring"""
					with open(lowercase__			,							encoding="utf-8"      ) as f:
										A			=					f.read().splitlines()[0]
					A			=					Text(encoding="utf-8"			,							features=Features({"image": Image()}      )      )
					A			=					text._generate_tables([[text_file_with_image]]      )
					A			=					pa.concat_tables([table for _, table in generator]      )
					assert pa_table.schema.field("image"      ).type == Image()()
					A			=					pa_table.to_pydict()["image"]
					assert generated_content == [{"path": image_file, "bytes": None}]
@pytest.mark.parametrize("sample_by"			,							["line", "paragraph", "document"]      )
def    __SCREAMING_SNAKE_CASE						(						lowercase__			,							lowercase__      ):
					"""simple docstring"""
					with open(lowercase__			,							encoding="utf-8"      ) as f:
										A			=					f.read()
					if sample_by == "line":
										A			=					expected_content.splitlines()
					elif sample_by == "paragraph":
										A			=					expected_content.split("\n\n"      )
					elif sample_by == "document":
										A			=					[expected_content]
					A			=					Text(sample_by=lowercase__			,							encoding="utf-8"			,							chunksize=100      )
					A			=					text._generate_tables([[text_file]]      )
					A			=					pa.concat_tables([table for _, table in generator]      ).to_pydict()["text"]
					assert generated_content == expected_content





