






"""simple docstring"""






import gc
import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 							lowerCAmelCase__				(			UpperCAmelCase__				,		UpperCAmelCase__				,		unittest.TestCase ):






      '''simple docstring'''

      __UpperCamelCase	 =  StableDiffusionSAGPipeline
      __UpperCamelCase	 =  TEXT_TO_IMAGE_PARAMS
      __UpperCamelCase	 =  TEXT_TO_IMAGE_BATCH_PARAMS
      __UpperCamelCase	 =  TEXT_TO_IMAGE_IMAGE_PARAMS
      __UpperCamelCase	 =  TEXT_TO_IMAGE_IMAGE_PARAMS
      __UpperCamelCase	 =  False
      def _SCREAMING_SNAKE_CASE (						self     :  str):




          '''simple docstring'''


          torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		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						,	)
          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		DDIMScheduler(
              beta_start=0.0_00_85						,	beta_end=0.0_12						,	beta_schedule='''scaled_linear'''						,	clip_sample=lowercase_						,	set_alpha_to_one=lowercase_						,	)
          torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_      :  Any    						=		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)
          SCREAMING_SNAKE_CASE_      :  str    						=		CLIPTextConfig(
              bos_token_id=0						,	eos_token_id=2						,	hidden_size=32						,	intermediate_size=37						,	layer_norm_eps=1e-05						,	num_attention_heads=4						,	num_hidden_layers=5						,	pad_token_id=1						,	vocab_size=1000						,	)
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		CLIPTextModel(lowercase_)
          SCREAMING_SNAKE_CASE_      :  int    						=		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

          SCREAMING_SNAKE_CASE_      :  str    						=		{
              '''unet''': unet,
              '''scheduler''': scheduler,
              '''vae''': vae,
              '''text_encoder''': text_encoder,
              '''tokenizer''': tokenizer,
              '''safety_checker''': None,
              '''feature_extractor''': None,
          }
          return components
      def _SCREAMING_SNAKE_CASE (						self     :  Any						,	lowercase_     :  Optional[Any]						,	lowercase_     :  Tuple=0):




          '''simple docstring'''


          if str(lowercase_).startswith('''mps'''):
              SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		torch.manual_seed(lowercase_)
          else:
              SCREAMING_SNAKE_CASE_      :  str    						=		torch.Generator(device=lowercase_).manual_seed(lowercase_)
          SCREAMING_SNAKE_CASE_      :  str    						=		{
              '''prompt''': '''.''',
              '''generator''': generator,
              '''num_inference_steps''': 2,
              '''guidance_scale''': 1.0,
              '''sag_scale''': 1.0,
              '''output_type''': '''numpy''',
          }
          return inputs



      def _SCREAMING_SNAKE_CASE (						self     :  Optional[int]):




          '''simple docstring'''


          super().test_inference_batch_single_identical(expected_max_diff=3e-3)




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






      '''simple docstring'''

      def _SCREAMING_SNAKE_CASE (						self     :  Union[str, Any]):




          '''simple docstring'''


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
      def _SCREAMING_SNAKE_CASE (						self     :  Optional[Any]):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		StableDiffusionSAGPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4''')
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		sag_pipe.to(lowercase_)
          sag_pipe.set_progress_bar_config(disable=lowercase_)

          SCREAMING_SNAKE_CASE_      :  Dict    						=		'''.'''
          SCREAMING_SNAKE_CASE_      :  int    						=		torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		sag_pipe(
              [prompt]						,	generator=lowercase_						,	guidance_scale=7.5						,	sag_scale=1.0						,	num_inference_steps=20						,	output_type='''np''')

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		output.images

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          SCREAMING_SNAKE_CASE_      :  Dict    						=		np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-2
      def _SCREAMING_SNAKE_CASE (						self     :  Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Tuple    						=		StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base''')
          SCREAMING_SNAKE_CASE_      :  Dict    						=		sag_pipe.to(lowercase_)
          sag_pipe.set_progress_bar_config(disable=lowercase_)

          SCREAMING_SNAKE_CASE_      :  Dict    						=		'''.'''
          SCREAMING_SNAKE_CASE_      :  str    						=		torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_      :  Union[str, Any]    						=		sag_pipe(
              [prompt]						,	generator=lowercase_						,	guidance_scale=7.5						,	sag_scale=1.0						,	num_inference_steps=20						,	output_type='''np''')

          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		output.images

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		image[0, -3:, -3:, -1]

          assert image.shape == (1, 512, 512, 3)
          SCREAMING_SNAKE_CASE_      :  List[str]    						=		np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71])

          assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-2



      def _SCREAMING_SNAKE_CASE (						self     :  Dict):




          '''simple docstring'''


          SCREAMING_SNAKE_CASE_      :  Optional[Any]    						=		StableDiffusionSAGPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base''')
          SCREAMING_SNAKE_CASE_      :  List[Any]    						=		sag_pipe.to(lowercase_)
          sag_pipe.set_progress_bar_config(disable=lowercase_)

          SCREAMING_SNAKE_CASE_      :  Optional[int]    						=		'''.'''
          SCREAMING_SNAKE_CASE_      :  Tuple    						=		torch.manual_seed(0)
          SCREAMING_SNAKE_CASE_      :  str    						=		sag_pipe(
              [prompt]						,	width=768						,	height=512						,	generator=lowercase_						,	guidance_scale=7.5						,	sag_scale=1.0						,	num_inference_steps=20						,	output_type='''np'''						,	)

          SCREAMING_SNAKE_CASE_      :  List[str]    						=		output.images

          assert image.shape == (1, 512, 768, 3)







"""simple docstring"""

import copy
import os

import cva
import numpy as np
from matplotlib import pyplot as plt




class        _UpperCAmelCase       :

 def __init__(     self		:       List[Any]				):
     snake_case_ :     List[str]				=	''''''
     snake_case_ :     Tuple				=	''''''
     snake_case_ :     int				=	[]
     snake_case_ :     Optional[int]				=	0
     snake_case_ :     Optional[Any]				=	256
     snake_case_ :     Tuple				=	0
     snake_case_ :     Tuple				=	0
     snake_case_ :     Optional[Any]				=	0
     snake_case_ :     Any				=	0

 def 						_snake_case					(     self		:       Optional[Any]  ,     lowercase_		:       List[Any]				):
     snake_case_ :     List[Any]				=	cva.imread(lowercase_  ,     0				)
     snake_case_ :     Tuple				=	copy.deepcopy(self.img				)
     snake_case_,	snake_case_,	snake_case_ :     List[Any]				=	plt.hist(self.img.ravel()  ,     256  ,     [0, 256]  ,     label='''x'''				)
     snake_case_ :     str				=	np.sum(lowercase_				)
     for i in range(len(lowercase_				)				):
         snake_case_ :     Optional[Any]				=	x[i] / self.k
         self.sk += prk
         snake_case_ :     Any				=	(self.L - 1) * self.sk
         if self.rem != 0:
             snake_case_ :     Dict				=	int(last % last				)
         snake_case_ :     Union[str, Any]				=	int(last + 1 if self.rem >= 0.5 else last				)
         self.last_list.append(lowercase_				)
         snake_case_ :     int				=	int(np.ma.count(self.img				) / self.img[1].size				)
         snake_case_ :     Tuple				=	self.img[1].size
     for i in range(self.number_of_cols				):
         for j in range(self.number_of_rows				):
             snake_case_ :     Union[str, Any]				=	self.img[j][i]
             if num != self.last_list[num]:
                 snake_case_ :     List[str]				=	self.last_list[num]
     cva.imwrite('''output_data/output.jpg'''  ,     self.img				)

 def 						_snake_case					(     self		:       Tuple				):
     plt.hist(self.img.ravel()  ,     256  ,     [0, 256]				)






 def 						_snake_case					(     self		:       int				):
     cva.imshow('''Output-Image'''  ,     self.img				)
     cva.imshow('''Input-Image'''  ,     self.original_image				)
     cva.waitKey(5000				)
     cva.destroyAllWindows()


if __name__ == "__main__":
     lowercase__    :  Any							    =	os.path.join(os.path.basename(__file__), '''image_data/input.jpg''')
     lowercase__    :  Any							    =	ConstantStretch()
     stretcher.stretch(file_path)
     stretcher.plot_histogram()
     stretcher.show_image()





