
import argparse
import os

import torch

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)


UpperCAmelCase	      = {
    '''sample_size''': 32,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 2,
    '''num_class_embeds''': 1000,
    '''block_out_channels''': [32, 64],
    '''attention_head_dim''': 8,
    '''down_block_types''': [
        '''ResnetDownsampleBlock2D''',
        '''AttnDownBlock2D''',
    ],
    '''up_block_types''': [
        '''AttnUpBlock2D''',
        '''ResnetUpsampleBlock2D''',
    ],
    '''resnet_time_scale_shift''': '''scale_shift''',
    '''upsample_type''': '''resnet''',
    '''downsample_type''': '''resnet''',
}

UpperCAmelCase	      = {
    '''sample_size''': 64,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 3,
    '''num_class_embeds''': 1000,
    '''block_out_channels''': [192, 192 * 2, 192 * 3, 192 * 4],
    '''attention_head_dim''': 64,
    '''down_block_types''': [
        '''ResnetDownsampleBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
    ],
    '''up_block_types''': [
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''ResnetUpsampleBlock2D''',
    ],
    '''resnet_time_scale_shift''': '''scale_shift''',
    '''upsample_type''': '''resnet''',
    '''downsample_type''': '''resnet''',
}

UpperCAmelCase	      = {
    '''sample_size''': 256,
    '''in_channels''': 3,
    '''out_channels''': 3,
    '''layers_per_block''': 2,
    '''num_class_embeds''': None,
    '''block_out_channels''': [256, 256, 256 * 2, 256 * 2, 256 * 4, 256 * 4],
    '''attention_head_dim''': 64,
    '''down_block_types''': [
        '''ResnetDownsampleBlock2D''',
        '''ResnetDownsampleBlock2D''',
        '''ResnetDownsampleBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
        '''AttnDownBlock2D''',
    ],
    '''up_block_types''': [
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''AttnUpBlock2D''',
        '''ResnetUpsampleBlock2D''',
        '''ResnetUpsampleBlock2D''',
        '''ResnetUpsampleBlock2D''',
    ],
    '''resnet_time_scale_shift''': '''default''',
    '''upsample_type''': '''resnet''',
    '''downsample_type''': '''resnet''',
}

UpperCAmelCase	      = {
    '''num_train_timesteps''': 40,
    '''sigma_min''': 0.002,
    '''sigma_max''': 80.0,
}

UpperCAmelCase	      = {
    '''num_train_timesteps''': 201,
    '''sigma_min''': 0.002,
    '''sigma_max''': 80.0,
}

UpperCAmelCase	      = {
    '''num_train_timesteps''': 151,
    '''sigma_min''': 0.002,
    '''sigma_max''': 80.0,
}






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	if isinstance(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
		return v
	if v.lower() in ("yes", "true", "t", "y", "1"):
		return True
	elif v.lower() in ("no", "false", "f", "n", "0"):
		return False
	else:
		raise argparse.ArgumentTypeError('boolean value expected'							)






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE=False							):
	lowercase       =    checkpoint[F'''{old_prefix}.in_layers.0.weight''']
	lowercase       =    checkpoint[F'''{old_prefix}.in_layers.0.bias''']
	lowercase       =    checkpoint[F'''{old_prefix}.in_layers.2.weight''']
	lowercase       =    checkpoint[F'''{old_prefix}.in_layers.2.bias''']
	lowercase       =    checkpoint[F'''{old_prefix}.emb_layers.1.weight''']
	lowercase       =    checkpoint[F'''{old_prefix}.emb_layers.1.bias''']
	lowercase       =    checkpoint[F'''{old_prefix}.out_layers.0.weight''']
	lowercase       =    checkpoint[F'''{old_prefix}.out_layers.0.bias''']
	lowercase       =    checkpoint[F'''{old_prefix}.out_layers.3.weight''']
	lowercase       =    checkpoint[F'''{old_prefix}.out_layers.3.bias''']

	if has_skip:
		lowercase       =    checkpoint[F'''{old_prefix}.skip_connection.weight''']
		lowercase       =    checkpoint[F'''{old_prefix}.skip_connection.bias''']

	return new_checkpoint






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE=None							):
	lowercase  ,					lowercase  ,					lowercase       =    checkpoint[F'''{old_prefix}.qkv.weight'''].chunk(3 ,							dim=0							)
	lowercase  ,					lowercase  ,					lowercase       =    checkpoint[F'''{old_prefix}.qkv.bias'''].chunk(3 ,							dim=0							)

	lowercase       =    checkpoint[F'''{old_prefix}.norm.weight''']
	lowercase       =    checkpoint[F'''{old_prefix}.norm.bias''']

	lowercase       =    weight_q.squeeze(-1							).squeeze(-1							)
	lowercase       =    bias_q.squeeze(-1							).squeeze(-1							)
	lowercase       =    weight_k.squeeze(-1							).squeeze(-1							)
	lowercase       =    bias_k.squeeze(-1							).squeeze(-1							)
	lowercase       =    weight_v.squeeze(-1							).squeeze(-1							)
	lowercase       =    bias_v.squeeze(-1							).squeeze(-1							)

	lowercase       =    (
	    checkpoint[F'''{old_prefix}.proj_out.weight'''].squeeze(-1							).squeeze(-1							)
	)
	lowercase       =    checkpoint[F'''{old_prefix}.proj_out.bias'''].squeeze(-1							).squeeze(-1							)

	return new_checkpoint






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    torch.load(__SCREAMING_SNAKE_CASE ,							map_location='cpu'							)
	lowercase       =    {}

	lowercase       =    checkpoint['time_embed.0.weight']
	lowercase       =    checkpoint['time_embed.0.bias']
	lowercase       =    checkpoint['time_embed.2.weight']
	lowercase       =    checkpoint['time_embed.2.bias']

	if unet_config["num_class_embeds"] is not None:
		lowercase       =    checkpoint['label_emb.weight']

	lowercase       =    checkpoint['input_blocks.0.0.weight']
	lowercase       =    checkpoint['input_blocks.0.0.bias']

	lowercase       =    unet_config['down_block_types']
	lowercase       =    unet_config['layers_per_block']
	lowercase       =    unet_config['attention_head_dim']
	lowercase       =    unet_config['block_out_channels']
	lowercase       =    1
	lowercase       =    channels_list[0]

	for i, layer_type in enumerate(__SCREAMING_SNAKE_CASE							):
		lowercase       =    channels_list[i]
		lowercase       =    current_channels != prev_channels
		if layer_type == "ResnetDownsampleBlock2D":
			for j in range(__SCREAMING_SNAKE_CASE							):
				lowercase       =    F'''down_blocks.{i}.resnets.{j}'''
				lowercase       =    F'''input_blocks.{current_layer}.0'''
				lowercase       =    True if j == 0 and downsample_block_has_skip else False
				lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							has_skip=__SCREAMING_SNAKE_CASE							)
				current_layer += 1

		elif layer_type == "AttnDownBlock2D":
			for j in range(__SCREAMING_SNAKE_CASE							):
				lowercase       =    F'''down_blocks.{i}.resnets.{j}'''
				lowercase       =    F'''input_blocks.{current_layer}.0'''
				lowercase       =    True if j == 0 and downsample_block_has_skip else False
				lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							has_skip=__SCREAMING_SNAKE_CASE							)
				lowercase       =    F'''down_blocks.{i}.attentions.{j}'''
				lowercase       =    F'''input_blocks.{current_layer}.1'''
				lowercase       =    convert_attention(
				    __SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
				current_layer += 1

		if i != len(__SCREAMING_SNAKE_CASE							) - 1:
			lowercase       =    F'''down_blocks.{i}.downsamplers.0'''
			lowercase       =    F'''input_blocks.{current_layer}.0'''
			lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
			current_layer += 1

		lowercase       =    current_channels

	# hardcoded the mid-block for now
	lowercase       =    'mid_block.resnets.0'
	lowercase       =    'middle_block.0'
	lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
	lowercase       =    'mid_block.attentions.0'
	lowercase       =    'middle_block.1'
	lowercase       =    convert_attention(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
	lowercase       =    'mid_block.resnets.1'
	lowercase       =    'middle_block.2'
	lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

	lowercase       =    0
	lowercase       =    unet_config['up_block_types']

	for i, layer_type in enumerate(__SCREAMING_SNAKE_CASE							):
		if layer_type == "ResnetUpsampleBlock2D":
			for j in range(layers_per_block + 1							):
				lowercase       =    F'''up_blocks.{i}.resnets.{j}'''
				lowercase       =    F'''output_blocks.{current_layer}.0'''
				lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							has_skip=__SCREAMING_SNAKE_CASE							)
				current_layer += 1

			if i != len(__SCREAMING_SNAKE_CASE							) - 1:
				lowercase       =    F'''up_blocks.{i}.upsamplers.0'''
				lowercase       =    F'''output_blocks.{current_layer-1}.1'''
				lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
		elif layer_type == "AttnUpBlock2D":
			for j in range(layers_per_block + 1							):
				lowercase       =    F'''up_blocks.{i}.resnets.{j}'''
				lowercase       =    F'''output_blocks.{current_layer}.0'''
				lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							has_skip=__SCREAMING_SNAKE_CASE							)
				lowercase       =    F'''up_blocks.{i}.attentions.{j}'''
				lowercase       =    F'''output_blocks.{current_layer}.1'''
				lowercase       =    convert_attention(
				    __SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
				current_layer += 1

			if i != len(__SCREAMING_SNAKE_CASE							) - 1:
				lowercase       =    F'''up_blocks.{i}.upsamplers.0'''
				lowercase       =    F'''output_blocks.{current_layer-1}.2'''
				lowercase       =    convert_resnet(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

	lowercase       =    checkpoint['out.0.weight']
	lowercase       =    checkpoint['out.0.bias']
	lowercase       =    checkpoint['out.2.weight']
	lowercase       =    checkpoint['out.2.bias']

	return new_checkpoint


if __name__ == "__main__":
	UpperCAmelCase	      = argparse.ArgumentParser()

	parser.add_argument('''--unet_path''', default=None, type=str, required=True, help='''Path to the unet.pt to convert.''')
	parser.add_argument(
	    '''--dump_path''', default=None, type=str, required=True, help='''Path to output the converted UNet model.'''
	)
	parser.add_argument('''--class_cond''', default=True, type=str, help='''Whether the model is class-conditional.''')

	UpperCAmelCase	      = parser.parse_args()
	UpperCAmelCase	      = strabool(args.class_cond)

	UpperCAmelCase	      = os.path.basename(args.unet_path)
	print(F"""Checkpoint: {ckpt_name}""")

	# Get U-Net config
	if "imagenet64" in ckpt_name:
		UpperCAmelCase	      = IMAGENET_64_UNET_CONFIG
	elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
		UpperCAmelCase	      = LSUN_256_UNET_CONFIG
	elif "test" in ckpt_name:
		UpperCAmelCase	      = TEST_UNET_CONFIG
	else:
		raise ValueError(F"""Checkpoint type {ckpt_name} is not currently supported.""")

	if not args.class_cond:
		UpperCAmelCase	      = None

	UpperCAmelCase	      = con_pt_to_diffuser(args.unet_path, unet_config)

	UpperCAmelCase	      = UNetaDModel(**unet_config)
	image_unet.load_state_dict(converted_unet_ckpt)

	# Get scheduler config
	if "cd" in ckpt_name or "test" in ckpt_name:
		UpperCAmelCase	      = CD_SCHEDULER_CONFIG
	elif "ct" in ckpt_name and "imagenet64" in ckpt_name:
		UpperCAmelCase	      = CT_IMAGENET_64_SCHEDULER_CONFIG
	elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
		UpperCAmelCase	      = CT_LSUN_256_SCHEDULER_CONFIG
	else:
		raise ValueError(F"""Checkpoint type {ckpt_name} is not currently supported.""")

	UpperCAmelCase	      = CMStochasticIterativeScheduler(**scheduler_config)

	UpperCAmelCase	      = ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler)
	consistency_model.save_pretrained(args.dump_path)



import sys
from collections import defaultdict







class        A_   :




	'''simple docstring'''


	def __init__(       self       ):
		lowercase       =    []





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case       ):
		return self.node_position[vertex]





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		lowercase       =    pos





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case       ):
		if start > size // 2 - 1:
			return
		else:
			if 2 * start + 2 >= size:
				lowercase       =    2 * start + 1
			else:
				if heap[2 * start + 1] < heap[2 * start + 2]:
					lowercase       =    2 * start + 1
				else:
					lowercase       =    2 * start + 2
			if heap[smallest_child] < heap[start]:
				lowercase  ,					lowercase       =    heap[smallest_child], positions[smallest_child]
				lowercase  ,					lowercase       =    (
				    heap[start],
				    positions[start],
				)
				lowercase  ,					lowercase       =    temp, tempa

				lowercase       =    self.get_position(positions[smallest_child]       )
				self.set_position(
				    positions[smallest_child]		,    self.get_position(positions[start]       )       )
				self.set_position(positions[start]		,    snake_case       )

				self.top_to_bottom(snake_case		,    snake_case		,    snake_case		,    snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case		,    snake_case		,    snake_case       ):
		lowercase       =    position[index]

		while index != 0:
			lowercase       =    int((index - 2) / 2       ) if index % 2 == 0 else int((index - 1) / 2       )

			if val < heap[parent]:
				lowercase       =    heap[parent]
				lowercase       =    position[parent]
				self.set_position(position[parent]		,    snake_case       )
			else:
				lowercase       =    val
				lowercase       =    temp
				self.set_position(snake_case		,    snake_case       )
				break
			lowercase       =    parent
		else:
			lowercase       =    val
			lowercase       =    temp
			self.set_position(snake_case		,    0       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		lowercase       =    len(snake_case       ) // 2 - 1
		for i in range(snake_case		,    -1		,    -1       ):
			self.top_to_bottom(snake_case		,    snake_case		,    len(snake_case       )		,    snake_case       )





	def 	SCREAMING_SNAKE_CASE__						(       self		,    snake_case		,    snake_case       ):
		lowercase       =    positions[0]
		lowercase       =    sys.maxsize
		self.top_to_bottom(snake_case		,    0		,    len(snake_case       )		,    snake_case       )
		return temp






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	lowercase       =    Heap()

	lowercase       =    [0] * len(__SCREAMING_SNAKE_CASE							)
	lowercase       =    [-1] * len(__SCREAMING_SNAKE_CASE							)  # Neighboring Tree Vertex of selected vertex
	# Minimum Distance of explored vertex with neighboring vertex of partial tree
	# formed in graph
	lowercase       =    []  # Heap of Distance of vertices from their neighboring vertex
	lowercase       =    []

	for vertex in range(len(__SCREAMING_SNAKE_CASE							)							):
		distance_tv.append(sys.maxsize							)
		positions.append(__SCREAMING_SNAKE_CASE							)
		heap.node_position.append(__SCREAMING_SNAKE_CASE							)

	lowercase       =    []
	lowercase       =    1
	lowercase       =    sys.maxsize
	for neighbor, distance in adjacency_list[0]:
		lowercase       =    0
		lowercase       =    distance
	heap.heapify(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)

	for _ in range(1 ,							len(__SCREAMING_SNAKE_CASE							)							):
		lowercase       =    heap.delete_minimum(__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
		if visited[vertex] == 0:
			tree_edges.append((nbr_tv[vertex], vertex)							)
			lowercase       =    1
			for neighbor, distance in adjacency_list[vertex]:
				if (
				    visited[neighbor] == 0
				    and distance < distance_tv[heap.get_position(__SCREAMING_SNAKE_CASE							)]
				):
					lowercase       =    distance
					heap.bottom_to_top(
					    __SCREAMING_SNAKE_CASE ,							heap.get_position(__SCREAMING_SNAKE_CASE							) ,							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							)
					lowercase       =    vertex
	return tree_edges


if __name__ == "__main__":  # pragma: no cover
	# < --------- Prims Algorithm --------- >
	UpperCAmelCase	      = int(input('''Enter number of edges: ''').strip())
	UpperCAmelCase	      = defaultdict(list)
	for _ in range(edges_number):
		UpperCAmelCase	      = [int(x) for x in input().strip().split()]
		adjacency_list[edge[0]].append([edge[1], edge[2]])
		adjacency_list[edge[1]].append([edge[0], edge[2]])
	print(prisms_algorithm(adjacency_list))

