




import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
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 ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()






def 						_snake_case					( _snake_case							: Optional[int]							)   -> List[Any]:






      '''simple docstring'''


      _A        =       [tensor.shape for tensor in tensor_list]
      return all(shape == shapes[0] for shape in shapes[1:]							)





class     lowercase_	(					__lowerCAmelCase      ,     __lowerCAmelCase      ,     __lowerCAmelCase      ,     unittest.TestCase							):



 '''simple docstring'''
 UpperCAmelCase   :					Tuple				  =     StableDiffusionLatentUpscalePipeline
 UpperCAmelCase   :					Tuple				  =     TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
     '''height''',
     '''width''',
     '''cross_attention_kwargs''',
     '''negative_prompt_embeds''',
     '''prompt_embeds''',
 }
 UpperCAmelCase   :					Optional[int]				  =     PipelineTesterMixin.required_optional_params - {'''num_images_per_prompt'''}
 UpperCAmelCase   :					Optional[Any]				  =     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
 UpperCAmelCase   :					Any				  =     frozenset(
     []							)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
 UpperCAmelCase   :					Tuple				  =     frozenset([]							)

 UpperCAmelCase   :					Tuple				  =     True

 @property
 def       lowerCAmelCase_							(			self	:	str		):
       _A        =       1
       _A        =       4
       _A        =       (16, 16)

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

 def       lowerCAmelCase_							(			self	:	List[str]		):
       torch.manual_seed(0		)
       _A        =       UNetaDConditionModel(
           act_fn='gelu'   ,							attention_head_dim=8   ,							norm_num_groups=_UpperCAmelCase   ,							block_out_channels=[32, 32, 64, 64]   ,							time_cond_proj_dim=160   ,							conv_in_kernel=1   ,							conv_out_kernel=1   ,							cross_attention_dim=32   ,							down_block_types=(
               'KDownBlock2D',
               'KCrossAttnDownBlock2D',
               'KCrossAttnDownBlock2D',
               'KCrossAttnDownBlock2D',
           )   ,							in_channels=8   ,							mid_block_type=_UpperCAmelCase   ,							only_cross_attention=_UpperCAmelCase   ,							out_channels=5   ,							resnet_time_scale_shift='scale_shift'   ,							time_embedding_type='fourier'   ,							timestep_post_act='gelu'   ,							up_block_types=('KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KCrossAttnUpBlock2D', 'KUpBlock2D')   ,							)
       _A        =       AutoencoderKL(
           block_out_channels=[32, 32, 64, 64]   ,							in_channels=3   ,							out_channels=3   ,							down_block_types=[
               'DownEncoderBlock2D',
               'DownEncoderBlock2D',
               'DownEncoderBlock2D',
               'DownEncoderBlock2D',
           ]   ,							up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D', 'UpDecoderBlock2D']   ,							latent_channels=4   ,							)
       _A        =       EulerDiscreteScheduler(prediction_type='sample'		)
       _A        =       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=1_000   ,							hidden_act='quick_gelu'   ,							projection_dim=512   ,							)
       _A        =       CLIPTextModel(_UpperCAmelCase		)
       _A        =       CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'		)

       _A        =       {
           'unet': model.eval(),
           'vae': vae.eval(),
           'scheduler': scheduler,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
       }

       return components

 def       lowerCAmelCase_							(			self	:	Tuple   ,							_UpperCAmelCase	:	Any   ,							_UpperCAmelCase	:	Optional[Any]=0		):
       if str(_UpperCAmelCase		).startswith('mps'		):
             _A        =       torch.manual_seed(_UpperCAmelCase		)
       else:
             _A        =       torch.Generator(device=_UpperCAmelCase		).manual_seed(_UpperCAmelCase		)
       _A        =       {
           'prompt': 'A painting of a squirrel eating a burger',
           'image': self.dummy_image.cpu(),
           'generator': generator,
           'num_inference_steps': 2,
           'output_type': 'numpy',
       }
       return inputs

 def       lowerCAmelCase_							(			self	:	Union[str, Any]		):
       _A        =       'cpu'

       _A        =       self.get_dummy_components()
       _A        =       self.pipeline_class(**_UpperCAmelCase		)
       pipe.to(_UpperCAmelCase		)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase		)

       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       pipe(**_UpperCAmelCase		).images
       _A        =       image[0, -3:, -3:, -1]

       self.assertEqual(image.shape   ,							(1, 256, 256, 3)		)
       _A        =       np.array(
           [0.4722_2412, 0.4192_1633, 0.4471_7434, 0.4687_4192, 0.4258_8258, 0.4615_0726, 0.467_7534, 0.4558_3832, 0.4857_9055]		)
       _A        =       np.abs(image_slice.flatten() - expected_slice		).max()
       self.assertLessEqual(_UpperCAmelCase   ,							1E-3		)

 def       lowerCAmelCase_							(			self	:	str		):
       super().test_attention_slicing_forward_pass(expected_max_diff=7E-3		)

 def       lowerCAmelCase_							(			self	:	int		):
       super().test_cpu_offload_forward_pass(expected_max_diff=3E-3		)

 def       lowerCAmelCase_							(			self	:	List[Any]		):
       super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3		)

 def       lowerCAmelCase_							(			self	:	Optional[int]		):
       super().test_inference_batch_single_identical(expected_max_diff=7E-3		)

 def       lowerCAmelCase_							(			self	:	List[Any]		):
       super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3E-3		)

 def       lowerCAmelCase_							(			self	:	Optional[Any]		):
       super().test_save_load_local(expected_max_difference=3E-3		)

 def       lowerCAmelCase_							(			self	:	List[Any]		):
       super().test_save_load_optional_components(expected_max_difference=3E-3		)



 def       lowerCAmelCase_							(			self	:	int		):
       _A        =       [
           'DDIMScheduler',
           'DDPMScheduler',
           'PNDMScheduler',
           'HeunDiscreteScheduler',
           'EulerAncestralDiscreteScheduler',
           'KDPM2DiscreteScheduler',
           'KDPM2AncestralDiscreteScheduler',
           'DPMSolverSDEScheduler',
       ]
       _A        =       self.get_dummy_components()
       _A        =       self.pipeline_class(**_UpperCAmelCase		)

       # make sure that PNDM does not need warm-up
       pipe.scheduler.register_to_config(skip_prk_steps=_UpperCAmelCase		)

       pipe.to(_UpperCAmelCase		)
       pipe.set_progress_bar_config(disable=_UpperCAmelCase		)
       _A        =       self.get_dummy_inputs(_UpperCAmelCase		)
       _A        =       2

       _A        =       []
       for scheduler_enum in KarrasDiffusionSchedulers:
             if scheduler_enum.name in skip_schedulers:
                   # no sigma schedulers are not supported
                   # no schedulers
                   continue

             _A        =       getattr(_UpperCAmelCase   ,							scheduler_enum.name		)
             _A        =       scheduler_cls.from_config(pipe.scheduler.config		)
             _A        =       pipe(**_UpperCAmelCase		)[0]
             outputs.append(_UpperCAmelCase		)

       assert check_same_shape(_UpperCAmelCase		)





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



 '''simple docstring'''

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

 def       lowerCAmelCase_							(			self	:	Optional[int]		):
       _A        =       torch.manual_seed(33		)

       _A        =       StableDiffusionPipeline.from_pretrained('CompVis/stable-diffusion-v1-4'   ,							torch_dtype=torch.floataa		)
       pipe.to('cuda'		)

       _A        =       StableDiffusionLatentUpscalePipeline.from_pretrained(
           'stabilityai/sd-x2-latent-upscaler'   ,							torch_dtype=torch.floataa		)
       upscaler.to('cuda'		)

       _A        =       'a photo of an astronaut high resolution, unreal engine, ultra realistic'

       _A        =       pipe(_UpperCAmelCase   ,							generator=_UpperCAmelCase   ,							output_type='latent'		).images

       _A        =       upscaler(
           prompt=_UpperCAmelCase   ,							image=_UpperCAmelCase   ,							num_inference_steps=20   ,							guidance_scale=0   ,							generator=_UpperCAmelCase   ,							output_type='np'   ,							).images[0]

       _A        =       load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy'		)
       assert np.abs((expected_image - image).mean()		) < 5E-2



 def       lowerCAmelCase_							(			self	:	str		):
       _A        =       torch.manual_seed(33		)

       _A        =       StableDiffusionLatentUpscalePipeline.from_pretrained(
           'stabilityai/sd-x2-latent-upscaler'   ,							torch_dtype=torch.floataa		)
       upscaler.to('cuda'		)

       _A        =       'the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas'

       _A        =       load_image(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png'		)

       _A        =       upscaler(
           prompt=_UpperCAmelCase   ,							image=_UpperCAmelCase   ,							num_inference_steps=20   ,							guidance_scale=0   ,							generator=_UpperCAmelCase   ,							output_type='np'   ,							).images[0]

       _A        =       load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy'		)
       assert np.abs((expected_image - image).max()		) < 5E-2







"""simple docstring"""

from collections import deque





class     lowercase_	:



 '''simple docstring'''

 def __init__(			self	:	int   ,							_UpperCAmelCase	:	str   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	int		):
       _A        =       process_name  # process name
       _A        =       arrival_time  # arrival time of the process
       # completion time of finished process or last interrupted time
       _A        =       arrival_time
       _A        =       burst_time  # remaining burst time
       _A        =       0  # total time of the process wait in ready queue
       _A        =       0  # time from arrival time to completion time





class     lowercase_	:



 '''simple docstring'''

 def __init__(			self	:	Optional[Any]   ,							_UpperCAmelCase	:	int   ,							_UpperCAmelCase	:	list[int]   ,							_UpperCAmelCase	:	deque[Process]   ,							_UpperCAmelCase	:	int   ,							):
       # total number of mlfq's queues
       _A        =       number_of_queues
       # time slice of queues that round robin algorithm applied
       _A        =       time_slices
       # unfinished process is in this ready_queue
       _A        =       queue
       # current time
       _A        =       current_time
       # finished process is in this sequence queue
       _A        =       deque()

 def       lowerCAmelCase_							(			self	:	Dict		):
       _A        =       []
       for i in range(len(self.finish_queue		)		):
             sequence.append(self.finish_queue[i].process_name		)
       return sequence

 def       lowerCAmelCase_							(			self	:	Optional[int]   ,							_UpperCAmelCase	:	list[Process]		):
       _A        =       []
       for i in range(len(_UpperCAmelCase		)		):
             waiting_times.append(queue[i].waiting_time		)
       return waiting_times

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	list[Process]		):
       _A        =       []
       for i in range(len(_UpperCAmelCase		)		):
             turnaround_times.append(queue[i].turnaround_time		)
       return turnaround_times

 def       lowerCAmelCase_							(			self	:	Optional[int]   ,							_UpperCAmelCase	:	list[Process]		):
       _A        =       []
       for i in range(len(_UpperCAmelCase		)		):
             completion_times.append(queue[i].stop_time		)
       return completion_times

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	deque[Process]		):
       return [q.burst_time for q in queue]

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	Process		):
       process.waiting_time += self.current_time - process.stop_time
       return process.waiting_time

 def       lowerCAmelCase_							(			self	:	int   ,							_UpperCAmelCase	:	deque[Process]		):
       _A        =       deque()  # sequence deque of finished process
       while len(_UpperCAmelCase		) != 0:
             _A        =       ready_queue.popleft()  # current process

             # if process's arrival time is later than current time, update current time
             if self.current_time < cp.arrival_time:
                   self.current_time += cp.arrival_time

             # update waiting time of current process
             self.update_waiting_time(_UpperCAmelCase		)
             # update current time
             self.current_time += cp.burst_time
             # finish the process and set the process's burst-time 0
             _A        =       0
             # set the process's turnaround time because it is finished
             _A        =       self.current_time - cp.arrival_time
             # set the completion time
             _A        =       self.current_time
             # add the process to queue that has finished queue
             finished.append(_UpperCAmelCase		)

       self.finish_queue.extend(_UpperCAmelCase		)  # add finished process to finish queue
       # FCFS will finish all remaining processes
       return finished

 def       lowerCAmelCase_							(			self	:	Dict   ,							_UpperCAmelCase	:	deque[Process]   ,							_UpperCAmelCase	:	int		):
       _A        =       deque()  # sequence deque of terminated process
       # just for 1 cycle and unfinished processes will go back to queue
       for _ in range(len(_UpperCAmelCase		)		):
             _A        =       ready_queue.popleft()  # current process

             # if process's arrival time is later than current time, update current time
             if self.current_time < cp.arrival_time:
                   self.current_time += cp.arrival_time

             # update waiting time of unfinished processes
             self.update_waiting_time(_UpperCAmelCase		)
             # if the burst time of process is bigger than time-slice
             if cp.burst_time > time_slice:
                   # use CPU for only time-slice
                   self.current_time += time_slice
                   # update remaining burst time
                   cp.burst_time -= time_slice
                   # update end point time
                   _A        =       self.current_time
                   # locate the process behind the queue because it is not finished
                   ready_queue.append(_UpperCAmelCase		)
             else:
                   # use CPU for remaining burst time
                   self.current_time += cp.burst_time
                   # set burst time 0 because the process is finished
                   _A        =       0
                   # set the finish time
                   _A        =       self.current_time
                   # update the process' turnaround time because it is finished
                   _A        =       self.current_time - cp.arrival_time
                   # add the process to queue that has finished queue
                   finished.append(_UpperCAmelCase		)

       self.finish_queue.extend(_UpperCAmelCase		)  # add finished process to finish queue
       # return finished processes queue and remaining processes queue
       return finished, ready_queue



 def       lowerCAmelCase_							(			self	:	str		):

       #  all queues except last one have round_robin algorithm
       for i in range(self.number_of_queues - 1		):
             _A       ,   _A        =       self.round_robin(
                 self.ready_queue   ,							self.time_slices[i]		)
       #  the last queue has first_come_first_served algorithm
       self.first_come_first_served(self.ready_queue		)

       return self.finish_queue


if __name__ == "__main__":
       import doctest

       a         =      Process('''P1''', 0, 53)
       a         =      Process('''P2''', 0, 17)
       a         =      Process('''P3''', 0, 68)
       a         =      Process('''P4''', 0, 24)
       a         =      3
       a         =      [17, 25]
       a         =      deque([Pa, Pa, Pa, Pa])

       if len(time_slices) != number_of_queues - 1:
              raise SystemExit(0)

       doctest.testmod(extraglobs={'''queue''': deque([Pa, Pa, Pa, Pa])})

       a         =      Process('''P1''', 0, 53)
       a         =      Process('''P2''', 0, 17)
       a         =      Process('''P3''', 0, 68)
       a         =      Process('''P4''', 0, 24)
       a         =      3
       a         =      [17, 25]
       a         =      deque([Pa, Pa, Pa, Pa])
       a         =      MLFQ(number_of_queues, time_slices, queue, 0)
       a         =      mlfq.multi_level_feedback_queue()

       # print total waiting times of processes(P1, P2, P3, P4)
       print(
           F'''waiting time:\
        \t\t\t{MLFQ.calculate_waiting_time(mlfq, [Pa, Pa, Pa, Pa])}'''
       )
       # print completion times of processes(P1, P2, P3, P4)
       print(
           F'''completion time:\
        \t\t{MLFQ.calculate_completion_time(mlfq, [Pa, Pa, Pa, Pa])}'''
       )
       # print total turnaround times of processes(P1, P2, P3, P4)
       print(
           F'''turnaround time:\
        \t\t{MLFQ.calculate_turnaround_time(mlfq, [Pa, Pa, Pa, Pa])}'''
       )
       # print sequence of finished processes
       print(
           F'''sequence of finished processes:\
        {mlfq.calculate_sequence_of_finish_queue()}'''
       )

