

"""simple docstring"""

import json
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
    import torch

    if is_vision_available():
        from transformers import OneFormerImageProcessor
        from transformers.models.oneformer.image_processing_oneformer import binary_mask_to_rle
        from transformers.models.oneformer.modeling_oneformer import OneFormerForUniversalSegmentationOutput

if is_vision_available():
    from PIL import Image



def       SCREAMING_SNAKE_CASE							( lowercase__					,					lowercase__="shi-labs/oneformer_demo"   )			->						List[Any]:
 with open(hf_hub_download(lowercase__					,					lowercase__					,					repo_type="dataset"   )					,					"r"   ) as f:
  lowerCAmelCase__   :				Optional[int] 			=				json.load(lowercase__   )
 lowerCAmelCase__   :				List[Any] 			=				{}
 lowerCAmelCase__   :				Optional[Any] 			=				[]
 lowerCAmelCase__   :				Any 			=				[]
 for key, info in class_info.items():
  lowerCAmelCase__   :				str 			=				info["name"]
  class_names.append(info["name"]   )
  if info["isthing"]:
   thing_ids.append(int(lowercase__   )   )
 lowerCAmelCase__   :				int 			=				thing_ids
 lowerCAmelCase__   :				Optional[Any] 			=				class_names
 return metadata






class 			__a					(						unittest.TestCase       ):





       """simple docstring"""


       def __init__(   self			,					snake_case			,					snake_case=7			,					snake_case=3			,					snake_case=30			,					snake_case=400			,					snake_case=None			,					snake_case=True			,					snake_case=True			,					snake_case=[0.5, 0.5, 0.5]			,					snake_case=[0.5, 0.5, 0.5]			,					snake_case=10			,					snake_case=False			,					snake_case=255			,					snake_case="shi-labs/oneformer_demo"			,					snake_case="ade20k_panoptic.json"			,					snake_case=10			,					):





        """simple docstring"""

        lowerCAmelCase__   :				int 			=				parent
        lowerCAmelCase__   :				Union[str, Any] 			=				batch_size
        lowerCAmelCase__   :				Optional[int] 			=				num_channels
        lowerCAmelCase__   :				int 			=				min_resolution
        lowerCAmelCase__   :				Dict 			=				max_resolution
        lowerCAmelCase__   :				str 			=				do_resize
        lowerCAmelCase__   :				Optional[Any] 			=				{"shortest_edge": 32, "longest_edge": 1_333} if size is None else size
        lowerCAmelCase__   :				Dict 			=				do_normalize
        lowerCAmelCase__   :				int 			=				image_mean
        lowerCAmelCase__   :				Union[str, Any] 			=				image_std
        lowerCAmelCase__   :				List[Any] 			=				class_info_file
        lowerCAmelCase__   :				str 			=				prepare_metadata(snake_case			,					snake_case   )
        lowerCAmelCase__   :				List[str] 			=				num_text
        lowerCAmelCase__   :				Any 			=				repo_path

        # for the post_process_functions
        lowerCAmelCase__   :				str 			=				2
        lowerCAmelCase__   :				str 			=				10
        lowerCAmelCase__   :				Tuple 			=				10
        lowerCAmelCase__   :				Optional[int] 			=				3
        lowerCAmelCase__   :				Tuple 			=				4
        lowerCAmelCase__   :				List[str] 			=				num_labels
        lowerCAmelCase__   :				Dict 			=				do_reduce_labels
        lowerCAmelCase__   :				Dict 			=				ignore_index






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        return {
            "do_resize": self.do_resize,
            "size": self.size,
            "do_normalize": self.do_normalize,
            "image_mean": self.image_mean,
            "image_std": self.image_std,
            "num_labels": self.num_labels,
            "do_reduce_labels": self.do_reduce_labels,
            "ignore_index": self.ignore_index,
            "class_info_file": self.class_info_file,
            "metadata": self.metadata,
            "num_text": self.num_text,
        }






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case			,					snake_case=False   ):





        """simple docstring"""

        if not batched:
         lowerCAmelCase__   :				Union[str, Any] 			=				image_inputs[0]
         if isinstance(snake_case			,					Image.Image   ):
          lowerCAmelCase__			,						lowerCAmelCase__   :				int 			=				image.size
         else:
          lowerCAmelCase__			,						lowerCAmelCase__   :				Dict 			=				image.shape[1], image.shape[2]
         if w < h:
          lowerCAmelCase__   :				Optional[Any] 			=				int(self.size["shortest_edge"] * h / w   )
          lowerCAmelCase__   :				str 			=				self.size["shortest_edge"]
         elif w > h:
          lowerCAmelCase__   :				Any 			=				self.size["shortest_edge"]
          lowerCAmelCase__   :				Any 			=				int(self.size["shortest_edge"] * w / h   )
         else:
          lowerCAmelCase__   :				str 			=				self.size["shortest_edge"]
          lowerCAmelCase__   :				List[str] 			=				self.size["shortest_edge"]

        else:
         lowerCAmelCase__   :				Optional[Any] 			=				[]
         for image in image_inputs:
          lowerCAmelCase__			,						lowerCAmelCase__   :				List[str] 			=				self.get_expected_values([image]   )
          expected_values.append((expected_height, expected_width)   )
         lowerCAmelCase__   :				Optional[Any] 			=				max(snake_case			,					key=lambda snake_case   : item[0]   )[0]
         lowerCAmelCase__   :				Optional[int] 			=				max(snake_case			,					key=lambda snake_case   : item[1]   )[1]

        return expected_height, expected_width






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        return OneFormerForUniversalSegmentationOutput(
            # +1 for null class
            class_queries_logits=torch.randn((self.batch_size, self.num_queries, self.num_classes + 1)   )			,					masks_queries_logits=torch.randn((self.batch_size, self.num_queries, self.height, self.width)   )			,					)






@require_torch
@require_vision
class 			__a					(						__magic_name__       ,	unittest.TestCase       ):





       """simple docstring"""


       __UpperCamelCase					:      List[str]    =						OneFormerImageProcessor if (is_vision_available() and is_torch_available()) else None
       # only for test_image_processing_common.test_image_proc_to_json_string
       __UpperCamelCase					:      Union[str, Any]    =						image_processing_class






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[str] 			=				OneFormerImageProcessorTester(self   )






       @property
       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        return self.image_processing_tester.prepare_image_processor_dict()






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Tuple 			=				self.image_processing_class(**self.image_processor_dict   )
        self.assertTrue(hasattr(snake_case			,					"image_mean"   )   )
        self.assertTrue(hasattr(snake_case			,					"image_std"   )   )
        self.assertTrue(hasattr(snake_case			,					"do_normalize"   )   )
        self.assertTrue(hasattr(snake_case			,					"do_resize"   )   )
        self.assertTrue(hasattr(snake_case			,					"size"   )   )
        self.assertTrue(hasattr(snake_case			,					"ignore_index"   )   )
        self.assertTrue(hasattr(snake_case			,					"class_info_file"   )   )
        self.assertTrue(hasattr(snake_case			,					"num_text"   )   )
        self.assertTrue(hasattr(snake_case			,					"repo_path"   )   )
        self.assertTrue(hasattr(snake_case			,					"metadata"   )   )
        self.assertTrue(hasattr(snake_case			,					"do_reduce_labels"   )   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        pass






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Dict 			=				self.image_processing_class(**self.image_processor_dict   )
        # create random PIL images
        lowerCAmelCase__   :				Any 			=				prepare_image_inputs(self.image_processing_tester			,					equal_resolution=snake_case   )
        for image in image_inputs:
         self.assertIsInstance(snake_case			,					Image.Image   )

        # Test not batched input
        lowerCAmelCase__   :				str 			=				image_processor(image_inputs[0]			,					["semantic"]			,					return_tensors="pt"   ).pixel_values

        lowerCAmelCase__			,						lowerCAmelCase__   :				Union[str, Any] 			=				self.image_processing_tester.get_expected_values(snake_case   )

        self.assertEqual(
            encoded_images.shape			,					(1, self.image_processing_tester.num_channels, expected_height, expected_width)			,					)

        # Test batched
        lowerCAmelCase__			,						lowerCAmelCase__   :				Optional[int] 			=				self.image_processing_tester.get_expected_values(snake_case			,					batched=snake_case   )

        lowerCAmelCase__   :				Tuple 			=				image_processor(
            snake_case			,					["semantic"] * len(snake_case   )			,					return_tensors="pt"   ).pixel_values
        self.assertEqual(
            encoded_images.shape			,					(
                self.image_processing_tester.batch_size,
                self.image_processing_tester.num_channels,
                expected_height,
                expected_width,
            )			,					)






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Tuple 			=				self.image_processing_class(**self.image_processor_dict   )
        # create random numpy tensors
        lowerCAmelCase__   :				Optional[int] 			=				prepare_image_inputs(self.image_processing_tester			,					equal_resolution=snake_case			,					numpify=snake_case   )
        for image in image_inputs:
         self.assertIsInstance(snake_case			,					np.ndarray   )

        # Test not batched input
        lowerCAmelCase__   :				Tuple 			=				image_processor(image_inputs[0]			,					["semantic"]			,					return_tensors="pt"   ).pixel_values

        lowerCAmelCase__			,						lowerCAmelCase__   :				Tuple 			=				self.image_processing_tester.get_expected_values(snake_case   )

        self.assertEqual(
            encoded_images.shape			,					(1, self.image_processing_tester.num_channels, expected_height, expected_width)			,					)

        # Test batched
        lowerCAmelCase__			,						lowerCAmelCase__   :				List[str] 			=				self.image_processing_tester.get_expected_values(snake_case			,					batched=snake_case   )

        lowerCAmelCase__   :				Union[str, Any] 			=				image_processor(
            snake_case			,					["semantic"] * len(snake_case   )			,					return_tensors="pt"   ).pixel_values
        self.assertEqual(
            encoded_images.shape			,					(
                self.image_processing_tester.batch_size,
                self.image_processing_tester.num_channels,
                expected_height,
                expected_width,
            )			,					)






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				self.image_processing_class(**self.image_processor_dict   )
        # create random PyTorch tensors
        lowerCAmelCase__   :				List[Any] 			=				prepare_image_inputs(self.image_processing_tester			,					equal_resolution=snake_case			,					torchify=snake_case   )
        for image in image_inputs:
         self.assertIsInstance(snake_case			,					torch.Tensor   )

        # Test not batched input
        lowerCAmelCase__   :				List[Any] 			=				image_processor(image_inputs[0]			,					["semantic"]			,					return_tensors="pt"   ).pixel_values

        lowerCAmelCase__			,						lowerCAmelCase__   :				List[Any] 			=				self.image_processing_tester.get_expected_values(snake_case   )

        self.assertEqual(
            encoded_images.shape			,					(1, self.image_processing_tester.num_channels, expected_height, expected_width)			,					)

        # Test batched
        lowerCAmelCase__			,						lowerCAmelCase__   :				Any 			=				self.image_processing_tester.get_expected_values(snake_case			,					batched=snake_case   )

        lowerCAmelCase__   :				str 			=				image_processor(
            snake_case			,					["semantic"] * len(snake_case   )			,					return_tensors="pt"   ).pixel_values
        self.assertEqual(
            encoded_images.shape			,					(
                self.image_processing_tester.batch_size,
                self.image_processing_tester.num_channels,
                expected_height,
                expected_width,
            )			,					)






       def 			SCREAMING_SNAKE_CASE_	(   self			,					snake_case=False			,					snake_case=False			,					snake_case="np"   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[Any] 			=				self.image_processing_class(**self.image_processor_dict   )
        # prepare image and target
        lowerCAmelCase__   :				Union[str, Any] 			=				self.image_processing_tester.num_labels
        lowerCAmelCase__   :				Any 			=				None
        lowerCAmelCase__   :				int 			=				None
        lowerCAmelCase__   :				Dict 			=				prepare_image_inputs(self.image_processing_tester			,					equal_resolution=snake_case   )
        if with_segmentation_maps:
         lowerCAmelCase__   :				Tuple 			=				num_labels
         if is_instance_map:
          lowerCAmelCase__   :				Optional[Any] 			=				list(range(snake_case   )   ) * 2
          lowerCAmelCase__   :				str 			=				dict(enumerate(snake_case   )   )
         lowerCAmelCase__   :				Optional[int] 			=				[
             np.random.randint(0			,					high * 2			,					(img.size[1], img.size[0])   ).astype(np.uinta   ) for img in image_inputs
         ]
         if segmentation_type == "pil":
          lowerCAmelCase__   :				Optional[Any] 			=				[Image.fromarray(snake_case   ) for annotation in annotations]

        lowerCAmelCase__   :				Any 			=				image_processor(
            snake_case			,					["semantic"] * len(snake_case   )			,					snake_case			,					return_tensors="pt"			,					instance_id_to_semantic_id=snake_case			,					pad_and_return_pixel_mask=snake_case			,					)

        return inputs






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        pass






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        def common(snake_case=False			,					snake_case=None   ):
         lowerCAmelCase__   :				Tuple 			=				self.comm_get_image_processor_inputs(
             with_segmentation_maps=snake_case			,					is_instance_map=snake_case			,					segmentation_type=snake_case   )

         lowerCAmelCase__   :				str 			=				inputs["mask_labels"]
         lowerCAmelCase__   :				Union[str, Any] 			=				inputs["class_labels"]
         lowerCAmelCase__   :				int 			=				inputs["pixel_values"]
         lowerCAmelCase__   :				List[Any] 			=				inputs["text_inputs"]

         # check the batch_size
         for mask_label, class_label, text_input in zip(snake_case			,					snake_case			,					snake_case   ):
          self.assertEqual(mask_label.shape[0]			,					class_label.shape[0]   )
          # this ensure padding has happened
          self.assertEqual(mask_label.shape[1:]			,					pixel_values.shape[2:]   )
          self.assertEqual(len(snake_case   )			,					self.image_processing_tester.num_text   )

        common()
        common(is_instance_map=snake_case   )
        common(is_instance_map=snake_case			,					segmentation_type="pil"   )
        common(is_instance_map=snake_case			,					segmentation_type="pil"   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				str 			=				np.zeros((20, 50)   )
        lowerCAmelCase__   :				Union[str, Any] 			=				1
        lowerCAmelCase__   :				Union[str, Any] 			=				1
        lowerCAmelCase__   :				Union[str, Any] 			=				1

        lowerCAmelCase__   :				str 			=				binary_mask_to_rle(snake_case   )
        self.assertEqual(len(snake_case   )			,					4   )
        self.assertEqual(rle[0]			,					21   )
        self.assertEqual(rle[1]			,					45   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Dict 			=				self.image_processing_class(
            num_labels=self.image_processing_tester.num_classes			,					max_seq_length=77			,					task_seq_length=77			,					class_info_file="ade20k_panoptic.json"			,					num_text=self.image_processing_tester.num_text			,					repo_path="shi-labs/oneformer_demo"			,					)
        lowerCAmelCase__   :				List[Any] 			=				self.image_processing_tester.get_fake_oneformer_outputs()

        lowerCAmelCase__   :				Tuple 			=				fature_extractor.post_process_semantic_segmentation(snake_case   )

        self.assertEqual(len(snake_case   )			,					self.image_processing_tester.batch_size   )
        self.assertEqual(
            segmentation[0].shape			,					(
                self.image_processing_tester.height,
                self.image_processing_tester.width,
            )			,					)

        lowerCAmelCase__   :				Union[str, Any] 			=				[(1, 4) for i in range(self.image_processing_tester.batch_size   )]
        lowerCAmelCase__   :				Optional[int] 			=				fature_extractor.post_process_semantic_segmentation(snake_case			,					target_sizes=snake_case   )

        self.assertEqual(segmentation[0].shape			,					target_sizes[0]   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				Optional[int] 			=				self.image_processing_class(
            num_labels=self.image_processing_tester.num_classes			,					max_seq_length=77			,					task_seq_length=77			,					class_info_file="ade20k_panoptic.json"			,					num_text=self.image_processing_tester.num_text			,					repo_path="shi-labs/oneformer_demo"			,					)
        lowerCAmelCase__   :				Optional[Any] 			=				self.image_processing_tester.get_fake_oneformer_outputs()
        lowerCAmelCase__   :				Dict 			=				image_processor.post_process_instance_segmentation(snake_case			,					threshold=0   )

        self.assertTrue(len(snake_case   ) == self.image_processing_tester.batch_size   )
        for el in segmentation:
         self.assertTrue("segmentation" in el   )
         self.assertTrue("segments_info" in el   )
         self.assertEqual(type(el["segments_info"]   )			,					snake_case   )
         self.assertEqual(
             el["segmentation"].shape			,					(self.image_processing_tester.height, self.image_processing_tester.width)   )






       def 			SCREAMING_SNAKE_CASE_	(   self   ):





        """simple docstring"""

        lowerCAmelCase__   :				List[Any] 			=				self.image_processing_class(
            num_labels=self.image_processing_tester.num_classes			,					max_seq_length=77			,					task_seq_length=77			,					class_info_file="ade20k_panoptic.json"			,					num_text=self.image_processing_tester.num_text			,					repo_path="shi-labs/oneformer_demo"			,					)
        lowerCAmelCase__   :				int 			=				self.image_processing_tester.get_fake_oneformer_outputs()
        lowerCAmelCase__   :				Union[str, Any] 			=				image_processor.post_process_panoptic_segmentation(snake_case			,					threshold=0   )

        self.assertTrue(len(snake_case   ) == self.image_processing_tester.batch_size   )
        for el in segmentation:
         self.assertTrue("segmentation" in el   )
         self.assertTrue("segments_info" in el   )
         self.assertEqual(type(el["segments_info"]   )			,					snake_case   )
         self.assertEqual(
             el["segmentation"].shape			,					(self.image_processing_tester.height, self.image_processing_tester.width)   )









"""simple docstring"""

import json
import os

import torch

from diffusers import UNetaDModel


os.makedirs("""hub/hopper-medium-v2/unet/hor32""", exist_ok=True)
os.makedirs("""hub/hopper-medium-v2/unet/hor128""", exist_ok=True)

os.makedirs("""hub/hopper-medium-v2/value_function""", exist_ok=True)



def       SCREAMING_SNAKE_CASE							( lowercase__   )			->						Optional[Any]:
 if hor == 1_2_8:
  lowerCAmelCase__   :				str 			=				("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D")
  lowerCAmelCase__   :				int 			=				(3_2, 1_2_8, 2_5_6)
  lowerCAmelCase__   :				Optional[Any] 			=				("UpResnetBlock1D", "UpResnetBlock1D")

 elif hor == 3_2:
  lowerCAmelCase__   :				Any 			=				("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D")
  lowerCAmelCase__   :				Dict 			=				(3_2, 6_4, 1_2_8, 2_5_6)
  lowerCAmelCase__   :				Optional[Any] 			=				("UpResnetBlock1D", "UpResnetBlock1D", "UpResnetBlock1D")
 lowerCAmelCase__   :				Tuple 			=				torch.load(F"""/Users/bglickenhaus/Documents/diffuser/temporal_unet-hopper-mediumv2-hor{hor}.torch"""   )
 lowerCAmelCase__   :				str 			=				model.state_dict()
 lowerCAmelCase__   :				Union[str, Any] 			=				{
     "down_block_types": down_block_types,
     "block_out_channels": block_out_channels,
     "up_block_types": up_block_types,
     "layers_per_block": 1,
     "use_timestep_embedding": True,
     "out_block_type": "OutConv1DBlock",
     "norm_num_groups": 8,
     "downsample_each_block": False,
     "in_channels": 1_4,
     "out_channels": 1_4,
     "extra_in_channels": 0,
     "time_embedding_type": "positional",
     "flip_sin_to_cos": False,
     "freq_shift": 1,
     "sample_size": 6_5_5_3_6,
     "mid_block_type": "MidResTemporalBlock1D",
     "act_fn": "mish",
 }
 lowerCAmelCase__   :				str 			=				UNetaDModel(**lowercase__   )
 print(F"""length of state dict: {len(state_dict.keys()   )}"""   )
 print(F"""length of value function dict: {len(hf_value_function.state_dict().keys()   )}"""   )
 lowerCAmelCase__   :				Dict 			=				dict(zip(model.state_dict().keys()					,					hf_value_function.state_dict().keys()   )   )
 for k, v in mapping.items():
  lowerCAmelCase__   :				Tuple 			=				state_dict.pop(lowercase__   )
 hf_value_function.load_state_dict(lowercase__   )

 torch.save(hf_value_function.state_dict()					,					F"""hub/hopper-medium-v2/unet/hor{hor}/diffusion_pytorch_model.bin"""   )
 with open(F"""hub/hopper-medium-v2/unet/hor{hor}/config.json"""					,					"w"   ) as f:
  json.dump(lowercase__					,					lowercase__   )



def       SCREAMING_SNAKE_CASE							( )			->						Any:
 lowerCAmelCase__   :				Dict 			=				{
     "in_channels": 1_4,
     "down_block_types": ("DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"),
     "up_block_types": (),
     "out_block_type": "ValueFunction",
     "mid_block_type": "ValueFunctionMidBlock1D",
     "block_out_channels": (3_2, 6_4, 1_2_8, 2_5_6),
     "layers_per_block": 1,
     "downsample_each_block": True,
     "sample_size": 6_5_5_3_6,
     "out_channels": 1_4,
     "extra_in_channels": 0,
     "time_embedding_type": "positional",
     "use_timestep_embedding": True,
     "flip_sin_to_cos": False,
     "freq_shift": 1,
     "norm_num_groups": 8,
     "act_fn": "mish",
 }

 lowerCAmelCase__   :				Any 			=				torch.load("/Users/bglickenhaus/Documents/diffuser/value_function-hopper-mediumv2-hor32.torch"   )
 lowerCAmelCase__   :				Optional[int] 			=				model
 lowerCAmelCase__   :				Any 			=				UNetaDModel(**lowercase__   )
 print(F"""length of state dict: {len(state_dict.keys()   )}"""   )
 print(F"""length of value function dict: {len(hf_value_function.state_dict().keys()   )}"""   )

 lowerCAmelCase__   :				str 			=				dict(zip(state_dict.keys()					,					hf_value_function.state_dict().keys()   )   )
 for k, v in mapping.items():
  lowerCAmelCase__   :				List[Any] 			=				state_dict.pop(lowercase__   )

 hf_value_function.load_state_dict(lowercase__   )

 torch.save(hf_value_function.state_dict()					,					"hub/hopper-medium-v2/value_function/diffusion_pytorch_model.bin"   )
 with open("hub/hopper-medium-v2/value_function/config.json"					,					"w"   ) as f:
  json.dump(lowercase__					,					lowercase__   )


if __name__ == "__main__":
    unet(3_2)
    # unet(128)
    value_function()






