






import random
import unittest

from torch.utils.data import BatchSampler, DataLoader, IterableDataset

from accelerate import Accelerator
from accelerate.data_loader import (
    BatchSamplerShard,
    DataLoaderDispatcher,
    DataLoaderShard,
    IterableDatasetShard,
    SkipBatchSampler,
    SkipDataLoader,
    skip_first_batches,
)







class  _lowercase   (   a__						):
						def __init__(     self				,			a=0.01				,			a=1_0_0_0   ):
							snake_case__  :    Tuple							       =p_stop
							snake_case__  :    str							       =max_length



						def __iter__(     self   ):
							snake_case__  :    str							       =0
							snake_case__  :    List[Any]							       =False
							while not stop and count < self.max_length:
								yield count
								count += 1
								snake_case__  :    Any							       =random.random() < self.p_stop







class  _lowercase   (   unittest.TestCase						):
						def 						lowercase__					(     self				,			a				,			a				,			a=False				,			a=True   ):
							snake_case__  :    Tuple							       =[
							    BatchSamplerShard(_A				,			2				,			_A				,			split_batches=_A				,			even_batches=_A   )
							    for i in range(2   )
							]
							snake_case__  :    Optional[int]							       =[list(_A   ) for batch_sampler_shard in batch_sampler_shards]
							if not split_batches:
								self.assertListEqual([len(_A   ) for shard in batch_sampler_shards]				,			[len(_A   ) for e in expected]   )
							self.assertListEqual(_A				,			_A   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Union[str, Any]							       =BatchSampler(range(2_4   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    int							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
							]
							self.check_batch_sampler_shards(_A				,			_A   )

							snake_case__  :    int							       =BatchSampler(range(2_4   )				,			batch_size=3				,			drop_last=_A   )
							# Expected shouldn't change
							self.check_batch_sampler_shards(_A				,			_A   )

							# Check the shards when the dataset is a round multiple of batch size but not total batch size.
							snake_case__  :    Union[str, Any]							       =BatchSampler(range(2_1   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    int							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [0, 1, 2]],
							]
							self.check_batch_sampler_shards(_A				,			_A   )

							snake_case__  :    Dict							       =BatchSampler(range(2_1   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    List[str]							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A   )

							# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
							# num_processes batch.
							snake_case__  :    Dict							       =BatchSampler(range(2_2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    List[Any]							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 0, 1]],
							]
							self.check_batch_sampler_shards(_A				,			_A   )

							snake_case__  :    int							       =BatchSampler(range(2_2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Optional[int]							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A   )

							# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
							# num_processes batch.
							snake_case__  :    List[str]							       =BatchSampler(range(2_0   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Optional[int]							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 0]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [1, 2, 3]],
							]
							self.check_batch_sampler_shards(_A				,			_A   )

							snake_case__  :    str							       =BatchSampler(range(2_0   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Dict							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A   )

							# Check the shards when the dataset is very small.
							snake_case__  :    int							       =BatchSampler(range(2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    int							       =[[[0, 1, 0]], [[1, 0, 1]]]
							self.check_batch_sampler_shards(_A				,			_A   )

							snake_case__  :    Union[str, Any]							       =BatchSampler(range(2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    List[str]							       =[[], []]
							self.check_batch_sampler_shards(_A				,			_A   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Optional[int]							       =BatchSampler(range(2_4   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Optional[int]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )

							snake_case__  :    List[Any]							       =BatchSampler(range(2_4   )				,			batch_size=4				,			drop_last=_A   )
							# Expected shouldn't change
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )

							# Check the shards when the dataset is not a round multiple of batch size.
							snake_case__  :    Tuple							       =BatchSampler(range(2_2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Optional[Any]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [0, 1]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )

							snake_case__  :    Optional[Any]							       =BatchSampler(range(2_2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Optional[int]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )

							# Check the shards when the dataset is not a round multiple of batch size or num_processes.
							snake_case__  :    Tuple							       =BatchSampler(range(2_1   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Union[str, Any]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 0]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [1, 2]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )

							snake_case__  :    int							       =BatchSampler(range(2_1   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    int							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )

							# Check the shards when the dataset is very small.
							snake_case__  :    Tuple							       =BatchSampler(range(2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    List[str]							       =[[[0, 1]], [[0, 1]]]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )

							snake_case__  :    List[str]							       =BatchSampler(range(2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Tuple							       =[[], []]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Union[str, Any]							       =BatchSampler(range(2_4   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Any							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1, 2_2, 2_3]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							snake_case__  :    Dict							       =BatchSampler(range(2_4   )				,			batch_size=3				,			drop_last=_A   )
							# Expected shouldn't change
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							# Check the shards when the dataset is a round multiple of batch size but not total batch size.
							snake_case__  :    Tuple							       =BatchSampler(range(2_1   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    int							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							snake_case__  :    Optional[int]							       =BatchSampler(range(2_1   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    str							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							# Check the shards when the dataset is not a round multiple of batch size but has a multiple of
							# num_processes batch.
							snake_case__  :    str							       =BatchSampler(range(2_2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Tuple							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9, 2_0]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7], [2_1]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							snake_case__  :    Union[str, Any]							       =BatchSampler(range(2_2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Dict							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							# Check the shards when the dataset is not a round multiple of batch size but and has not a multiple of
							# num_processes batch.
							snake_case__  :    Tuple							       =BatchSampler(range(2_0   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    str							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4], [1_8, 1_9]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							snake_case__  :    List[Any]							       =BatchSampler(range(2_0   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Dict							       =[
							    [[0, 1, 2], [6, 7, 8], [1_2, 1_3, 1_4]],
							    [[3, 4, 5], [9, 1_0, 1_1], [1_5, 1_6, 1_7]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							# Check the shards when the dataset is very small.
							snake_case__  :    int							       =BatchSampler(range(2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    Dict							       =[[[0, 1]], []]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )

							snake_case__  :    str							       =BatchSampler(range(2   )				,			batch_size=3				,			drop_last=_A   )
							snake_case__  :    str							       =[[], []]
							self.check_batch_sampler_shards(_A				,			_A				,			even_batches=_A   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Optional[int]							       =BatchSampler(range(2_4   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Optional[int]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9], [2_2, 2_3]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )

							snake_case__  :    str							       =BatchSampler(range(2_4   )				,			batch_size=4				,			drop_last=_A   )
							# Expected shouldn't change
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )

							# Check the shards when the dataset is not a round multiple of batch size.
							snake_case__  :    List[str]							       =BatchSampler(range(2_2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Optional[Any]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0, 2_1]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )

							snake_case__  :    Union[str, Any]							       =BatchSampler(range(2_2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    List[Any]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )

							# Check the shards when the dataset is not a round multiple of batch size or num_processes.
							snake_case__  :    int							       =BatchSampler(range(2_1   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Any							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7], [2_0]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )

							snake_case__  :    Any							       =BatchSampler(range(2_1   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Optional[int]							       =[
							    [[0, 1], [4, 5], [8, 9], [1_2, 1_3], [1_6, 1_7]],
							    [[2, 3], [6, 7], [1_0, 1_1], [1_4, 1_5], [1_8, 1_9]],
							]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )

							# Check the shards when the dataset is very small.
							snake_case__  :    Dict							       =BatchSampler(range(2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    int							       =[[[0, 1]], []]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )

							snake_case__  :    Dict							       =BatchSampler(range(2   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    Tuple							       =[[], []]
							self.check_batch_sampler_shards(_A				,			_A				,			split_batches=_A				,			even_batches=_A   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Tuple							       =[[0, 1, 2], [3, 4], [5, 6, 7, 8], [9, 1_0, 1_1], [1_2, 1_3]]
							snake_case__  :    Tuple							       =[BatchSamplerShard(_A				,			2				,			_A				,			even_batches=_A   ) for i in range(2   )]

							self.assertEqual(len(batch_sampler_shards[0]   )				,			3   )
							self.assertEqual(len(batch_sampler_shards[1]   )				,			2   )

							self.assertListEqual(list(batch_sampler_shards[0]   )				,			[[0, 1, 2], [5, 6, 7, 8], [1_2, 1_3]]   )
							self.assertListEqual(list(batch_sampler_shards[1]   )				,			[[3, 4], [9, 1_0, 1_1]]   )
						def 						lowercase__					(     self				,			a				,			a				,			a				,			a=False				,			a=2				,			a=False   ):
							random.seed(_A   )
							snake_case__  :    Tuple							       =list(_A   )

							snake_case__  :    List[str]							       =[
							    IterableDatasetShard(
							        _A				,			batch_size=_A				,			drop_last=_A				,			num_processes=_A				,			process_index=_A				,			split_batches=_A				,			)
							    for i in range(_A   )
							]
							snake_case__  :    Any							       =[]
							for iterable_dataset_shard in iterable_dataset_shards:
								# Since our random iterable dataset will be... random... we need to use a seed to get reproducible results.
								random.seed(_A   )
								iterable_dataset_lists.append(list(_A   )   )

							snake_case__  :    Union[str, Any]							       =batch_size // num_processes if split_batches else batch_size
							# All iterable dataset shard should have the same length, a round multiple of shard_batch_size
							snake_case__  :    Dict							       =iterable_dataset_lists[0]
							for l in iterable_dataset_lists[1:]:
								self.assertEqual(len(_A   )				,			len(_A   )   )
								self.assertTrue(len(_A   ) % shard_batch_size == 0   )

							snake_case__  :    Union[str, Any]							       =[]
							for idx in range(0				,			len(_A   )				,			_A   ):
								for l in iterable_dataset_lists:
									observed += l[idx : idx + shard_batch_size]

							if not drop_last:
								while len(_A   ) < len(_A   ):
									reference += reference
							self.assertListEqual(_A				,			reference[: len(_A   )]   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Dict							       =4_2
							snake_case__  :    Union[str, Any]							       =RandomIterableDataset()

							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )
							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )
							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )
							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )

							# Edge case with a very small dataset
							snake_case__  :    Optional[int]							       =RandomIterableDataset(max_length=2   )

							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )
							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )
							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )
							self.check_iterable_dataset_shards(_A				,			_A				,			batch_size=4				,			drop_last=_A				,			split_batches=_A   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Dict							       =BatchSampler(range(1_6   )				,			batch_size=4				,			drop_last=_A   )
							snake_case__  :    int							       =SkipBatchSampler(_A				,			2   )
							self.assertListEqual(list(_A   )				,			[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]   )
						def 						lowercase__					(     self   ):
							snake_case__  :    List[str]							       =SkipDataLoader(list(range(1_6   )   )				,			batch_size=4				,			skip_batches=2   )
							self.assertListEqual([t.tolist() for t in dataloader]				,			[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]   )
						def 						lowercase__					(     self   ):
							snake_case__  :    Optional[int]							       =DataLoader(list(range(1_6   )   )				,			batch_size=4   )
							snake_case__  :    int							       =skip_first_batches(_A				,			num_batches=2   )
							self.assertListEqual([t.tolist() for t in new_dataloader]				,			[[8, 9, 1_0, 1_1], [1_2, 1_3, 1_4, 1_5]]   )
						def 						lowercase__					(     self   ):
							snake_case__  :    List[str]							       =DataLoaderShard(list(range(1_6   )   )				,			batch_size=4   )
							for idx, _ in enumerate(_A   ):
								self.assertEqual(dataloader.end_of_dataloader				,			idx == 3   )

							# Test it also works on the second iteration
							for idx, _ in enumerate(_A   ):
								self.assertEqual(dataloader.end_of_dataloader				,			idx == 3   )



						def 						lowercase__					(     self   ):
							Accelerator()
							snake_case__  :    Tuple							       =DataLoaderDispatcher(range(1_6   )				,			batch_size=4   )
							for idx, _ in enumerate(_A   ):
								self.assertEqual(dataloader.end_of_dataloader				,			idx == 3   )

							# Test it also works on the second iteration
							for idx, _ in enumerate(_A   ):
								self.assertEqual(dataloader.end_of_dataloader				,			idx == 3   )













import argparse

import torch

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import download_from_original_stable_diffusion_ckpt


if __name__ == "__main__":
       __lowerCamelCase	:		List[str]										=   argparse.ArgumentParser()

       parser.add_argument(
           """--checkpoint_path""", default=None, type=str, required=True, help="""Path to the checkpoint to convert."""
       )
       # !wget https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml
       parser.add_argument(
           """--original_config_file""",
           default=None,
           type=str,
           help="""The YAML config file corresponding to the original architecture.""",
       )
       parser.add_argument(
           """--num_in_channels""",
           default=None,
           type=int,
           help="""The number of input channels. If `None` number of input channels will be automatically inferred.""",
       )
       parser.add_argument(
           """--scheduler_type""",
           default="""pndm""",
           type=str,
           help="""Type of scheduler to use. Should be one of ['pndm', 'lms', 'ddim', 'euler', 'euler-ancestral', 'dpm']""",
       )
       parser.add_argument(
           """--pipeline_type""",
           default=None,
           type=str,
           help=(
               """The pipeline type. One of 'FrozenOpenCLIPEmbedder', 'FrozenCLIPEmbedder', 'PaintByExample'"""
               """. If `None` pipeline will be automatically inferred."""
           ),
       )
       parser.add_argument(
           """--image_size""",
           default=None,
           type=int,
           help=(
               """The image size that the model was trained on. Use 512 for Stable Diffusion v1.X and Stable Siffusion v2"""
               """ Base. Use 768 for Stable Diffusion v2."""
           ),
       )
       parser.add_argument(
           """--prediction_type""",
           default=None,
           type=str,
           help=(
               """The prediction type that the model was trained on. Use 'epsilon' for Stable Diffusion v1.X and Stable"""
               """ Diffusion v2 Base. Use 'v_prediction' for Stable Diffusion v2."""
           ),
       )
       parser.add_argument(
           """--extract_ema""",
           action="""store_true""",
           help=(
               """Only relevant for checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights"""
               """ or not. Defaults to `False`. Add `--extract_ema` to extract the EMA weights. EMA weights usually yield"""
               """ higher quality images for inference. Non-EMA weights are usually better to continue fine-tuning."""
           ),
       )
       parser.add_argument(
           """--upcast_attention""",
           action="""store_true""",
           help=(
               """Whether the attention computation should always be upcasted. This is necessary when running stable"""
               """ diffusion 2.1."""
           ),
       )
       parser.add_argument(
           """--from_safetensors""",
           action="""store_true""",
           help="""If `--checkpoint_path` is in `safetensors` format, load checkpoint with safetensors instead of PyTorch.""",
       )
       parser.add_argument(
           """--to_safetensors""",
           action="""store_true""",
           help="""Whether to store pipeline in safetensors format or not.""",
       )
       parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")
       parser.add_argument("""--device""", type=str, help="""Device to use (e.g. cpu, cuda:0, cuda:1, etc.)""")
       parser.add_argument(
           """--stable_unclip""",
           type=str,
           default=None,
           required=False,
           help="""Set if this is a stable unCLIP model. One of 'txt2img' or 'img2img'.""",
       )
       parser.add_argument(
           """--stable_unclip_prior""",
           type=str,
           default=None,
           required=False,
           help="""Set if this is a stable unCLIP txt2img model. Selects which prior to use. If `--stable_unclip` is set to `txt2img`, the karlo prior (https://huggingface.co/kakaobrain/karlo-v1-alpha/tree/main/prior) is selected by default.""",
       )
       parser.add_argument(
           """--clip_stats_path""",
           type=str,
           help="""Path to the clip stats file. Only required if the stable unclip model's config specifies `model.params.noise_aug_config.params.clip_stats_path`.""",
           required=False,
       )
       parser.add_argument(
           """--controlnet""", action="""store_true""", default=None, help="""Set flag if this is a controlnet checkpoint."""
       )
       parser.add_argument("""--half""", action="""store_true""", help="""Save weights in half precision.""")
       parser.add_argument(
           """--vae_path""",
           type=str,
           default=None,
           required=False,
           help="""Set to a path, hub id to an already converted vae to not convert it again.""",
       )
       __lowerCamelCase	:		Optional[Any]										=   parser.parse_args()

       __lowerCamelCase	:		Optional[int]										=   download_from_original_stable_diffusion_ckpt(
           checkpoint_path=args.checkpoint_path,
           original_config_file=args.original_config_file,
           image_size=args.image_size,
           prediction_type=args.prediction_type,
           model_type=args.pipeline_type,
           extract_ema=args.extract_ema,
           scheduler_type=args.scheduler_type,
           num_in_channels=args.num_in_channels,
           upcast_attention=args.upcast_attention,
           from_safetensors=args.from_safetensors,
           device=args.device,
           stable_unclip=args.stable_unclip,
           stable_unclip_prior=args.stable_unclip_prior,
           clip_stats_path=args.clip_stats_path,
           controlnet=args.controlnet,
           vae_path=args.vae_path,
       )

       if args.half:
              pipe.to(torch_dtype=torch.floataa)

       if args.controlnet:
              # only save the controlnet model
              pipe.controlnet.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)
       else:
              pipe.save_pretrained(args.dump_path, safe_serialization=args.to_safetensors)





