





from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch
import torch.nn as nn

from ..utils import BaseOutput, is_torch_version, randn_tensor
from .attention_processor import SpatialNorm
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block

@dataclass
class    lowerCAmelCase							(					__UpperCamelCase			):
 UpperCAmelCase__													=		42


class    lowerCAmelCase							(					nn.Module			):
 def __init__(      self  :      Tuple			,     UpperCAmelCase  :      Dict=3			,     UpperCAmelCase  :      Union[str, Any]=3			,     UpperCAmelCase  :      Optional[Any]=("DownEncoderBlock2D",)			,     UpperCAmelCase  :      List[str]=(64,)			,     UpperCAmelCase  :      Any=2			,     UpperCAmelCase  :      List[Any]=32			,     UpperCAmelCase  :      List[str]="silu"			,     UpperCAmelCase  :      Tuple=True			,     )			->	Tuple:
     super().__init__()
     lowerCamelCase__    :      Optional[Any]					 =		layers_per_block

     lowerCamelCase__    :      Dict					 =		torch.nn.Convad(
         UpperCAmelCase			,     block_out_channels[0]			,     kernel_size=3			,     stride=1			,     padding=1			,     )

     lowerCamelCase__    :      Tuple					 =		None
     lowerCamelCase__    :      Union[str, Any]					 =		nn.ModuleList([]	)

     # down
     lowerCamelCase__    :      Tuple					 =		block_out_channels[0]
     for i, down_block_type in enumerate(UpperCAmelCase	):
         lowerCamelCase__    :      List[Any]					 =		output_channel
         lowerCamelCase__    :      Dict					 =		block_out_channels[i]
         lowerCamelCase__    :      int					 =		i == len(UpperCAmelCase	) - 1

         lowerCamelCase__    :      int					 =		get_down_block(
             UpperCAmelCase			,     num_layers=self.layers_per_block			,     in_channels=UpperCAmelCase			,     out_channels=UpperCAmelCase			,     add_downsample=not is_final_block			,     resnet_eps=1e-6			,     downsample_padding=0			,     resnet_act_fn=UpperCAmelCase			,     resnet_groups=UpperCAmelCase			,     attention_head_dim=UpperCAmelCase			,     temb_channels=UpperCAmelCase			,     )
         self.down_blocks.append(UpperCAmelCase	)

     # mid
     lowerCamelCase__    :      Tuple					 =		UNetMidBlockaD(
         in_channels=block_out_channels[-1]			,     resnet_eps=1e-6			,     resnet_act_fn=UpperCAmelCase			,     output_scale_factor=1			,     resnet_time_scale_shift='default'			,     attention_head_dim=block_out_channels[-1]			,     resnet_groups=UpperCAmelCase			,     temb_channels=UpperCAmelCase			,     )

     # out
     lowerCamelCase__    :      Optional[Any]					 =		nn.GroupNorm(num_channels=block_out_channels[-1]			,     num_groups=UpperCAmelCase			,     eps=1e-6	)
     lowerCamelCase__    :      List[Any]					 =		nn.SiLU()

     lowerCamelCase__    :      List[str]					 =		2 * out_channels if double_z else out_channels
     lowerCamelCase__    :      Any					 =		nn.Convad(block_out_channels[-1]			,     UpperCAmelCase			,     3			,     padding=1	)

     lowerCamelCase__    :      List[Any]					 =		False







 def      A_	(      self  :      Dict			,     UpperCAmelCase  :      Optional[int]	)			->	Union[str, Any]:
     lowerCamelCase__    :      Dict					 =		x
     lowerCamelCase__    :      Dict					 =		self.conv_in(UpperCAmelCase	)

     if self.training and self.gradient_checkpointing:

         def create_custom_forward(UpperCAmelCase  :      Dict	):
             def custom_forward(*UpperCAmelCase  :      Any	):
                 return module(*UpperCAmelCase	)

             return custom_forward

         # down
         if is_torch_version('>='			,     '1.11.0'	):
             for down_block in self.down_blocks:
                 lowerCamelCase__    :      Optional[Any]					 =		torch.utils.checkpoint.checkpoint(
                     create_custom_forward(UpperCAmelCase	)			,     UpperCAmelCase			,     use_reentrant=UpperCAmelCase	)
             # middle
             lowerCamelCase__    :      Tuple					 =		torch.utils.checkpoint.checkpoint(
                 create_custom_forward(self.mid_block	)			,     UpperCAmelCase			,     use_reentrant=UpperCAmelCase	)
         else:
             for down_block in self.down_blocks:
                 lowerCamelCase__    :      str					 =		torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase	)			,     UpperCAmelCase	)
             # middle
             lowerCamelCase__    :      Dict					 =		torch.utils.checkpoint.checkpoint(create_custom_forward(self.mid_block	)			,     UpperCAmelCase	)

     else:
         # down
         for down_block in self.down_blocks:
             lowerCamelCase__    :      Optional[Any]					 =		down_block(UpperCAmelCase	)

         # middle
         lowerCamelCase__    :      List[str]					 =		self.mid_block(UpperCAmelCase	)

     # post-process
     lowerCamelCase__    :      str					 =		self.conv_norm_out(UpperCAmelCase	)
     lowerCamelCase__    :      List[Any]					 =		self.conv_act(UpperCAmelCase	)
     lowerCamelCase__    :      Union[str, Any]					 =		self.conv_out(UpperCAmelCase	)

     return sample


class    lowerCAmelCase							(					nn.Module			):
 def __init__(      self  :      Any			,     UpperCAmelCase  :      Optional[int]=3			,     UpperCAmelCase  :      List[Any]=3			,     UpperCAmelCase  :      int=("UpDecoderBlock2D",)			,     UpperCAmelCase  :      Optional[Any]=(64,)			,     UpperCAmelCase  :      str=2			,     UpperCAmelCase  :      Union[str, Any]=32			,     UpperCAmelCase  :      int="silu"			,     UpperCAmelCase  :      List[str]="group"			,     )			->	Any:
     super().__init__()
     lowerCamelCase__    :      List[str]					 =		layers_per_block

     lowerCamelCase__    :      Optional[Any]					 =		nn.Convad(
         UpperCAmelCase			,     block_out_channels[-1]			,     kernel_size=3			,     stride=1			,     padding=1			,     )

     lowerCamelCase__    :      Dict					 =		None
     lowerCamelCase__    :      List[str]					 =		nn.ModuleList([]	)

     lowerCamelCase__    :      Dict					 =		in_channels if norm_type == 'spatial' else None

     # mid
     lowerCamelCase__    :      List[Any]					 =		UNetMidBlockaD(
         in_channels=block_out_channels[-1]			,     resnet_eps=1e-6			,     resnet_act_fn=UpperCAmelCase			,     output_scale_factor=1			,     resnet_time_scale_shift='default' if norm_type == 'group' else norm_type			,     attention_head_dim=block_out_channels[-1]			,     resnet_groups=UpperCAmelCase			,     temb_channels=UpperCAmelCase			,     )

     # up
     lowerCamelCase__    :      Tuple					 =		list(reversed(UpperCAmelCase	)	)
     lowerCamelCase__    :      Dict					 =		reversed_block_out_channels[0]
     for i, up_block_type in enumerate(UpperCAmelCase	):
         lowerCamelCase__    :      Dict					 =		output_channel
         lowerCamelCase__    :      List[Any]					 =		reversed_block_out_channels[i]

         lowerCamelCase__    :      Union[str, Any]					 =		i == len(UpperCAmelCase	) - 1

         lowerCamelCase__    :      Union[str, Any]					 =		get_up_block(
             UpperCAmelCase			,     num_layers=self.layers_per_block + 1			,     in_channels=UpperCAmelCase			,     out_channels=UpperCAmelCase			,     prev_output_channel=UpperCAmelCase			,     add_upsample=not is_final_block			,     resnet_eps=1e-6			,     resnet_act_fn=UpperCAmelCase			,     resnet_groups=UpperCAmelCase			,     attention_head_dim=UpperCAmelCase			,     temb_channels=UpperCAmelCase			,     resnet_time_scale_shift=UpperCAmelCase			,     )
         self.up_blocks.append(UpperCAmelCase	)
         lowerCamelCase__    :      str					 =		output_channel

     # out
     if norm_type == "spatial":
         lowerCamelCase__    :      Any					 =		SpatialNorm(block_out_channels[0]			,     UpperCAmelCase	)
     else:
         lowerCamelCase__    :      Optional[Any]					 =		nn.GroupNorm(num_channels=block_out_channels[0]			,     num_groups=UpperCAmelCase			,     eps=1e-6	)
     lowerCamelCase__    :      Tuple					 =		nn.SiLU()
     lowerCamelCase__    :      Union[str, Any]					 =		nn.Convad(block_out_channels[0]			,     UpperCAmelCase			,     3			,     padding=1	)

     lowerCamelCase__    :      Dict					 =		False







 def      A_	(      self  :      Tuple			,     UpperCAmelCase  :      Tuple			,     UpperCAmelCase  :      Any=None	)			->	List[str]:
     lowerCamelCase__    :      Any					 =		z
     lowerCamelCase__    :      Union[str, Any]					 =		self.conv_in(UpperCAmelCase	)

     lowerCamelCase__    :      Optional[int]					 =		next(iter(self.up_blocks.parameters()	)	).dtype
     if self.training and self.gradient_checkpointing:

         def create_custom_forward(UpperCAmelCase  :      Any	):
             def custom_forward(*UpperCAmelCase  :      Any	):
                 return module(*UpperCAmelCase	)

             return custom_forward

         if is_torch_version('>='			,     '1.11.0'	):
             # middle
             lowerCamelCase__    :      str					 =		torch.utils.checkpoint.checkpoint(
                 create_custom_forward(self.mid_block	)			,     UpperCAmelCase			,     UpperCAmelCase			,     use_reentrant=UpperCAmelCase	)
             lowerCamelCase__    :      Tuple					 =		sample.to(UpperCAmelCase	)

             # up
             for up_block in self.up_blocks:
                 lowerCamelCase__    :      str					 =		torch.utils.checkpoint.checkpoint(
                     create_custom_forward(UpperCAmelCase	)			,     UpperCAmelCase			,     UpperCAmelCase			,     use_reentrant=UpperCAmelCase	)
         else:
             # middle
             lowerCamelCase__    :      str					 =		torch.utils.checkpoint.checkpoint(
                 create_custom_forward(self.mid_block	)			,     UpperCAmelCase			,     UpperCAmelCase	)
             lowerCamelCase__    :      Dict					 =		sample.to(UpperCAmelCase	)

             # up
             for up_block in self.up_blocks:
                 lowerCamelCase__    :      Tuple					 =		torch.utils.checkpoint.checkpoint(create_custom_forward(UpperCAmelCase	)			,     UpperCAmelCase			,     UpperCAmelCase	)
     else:
         # middle
         lowerCamelCase__    :      Any					 =		self.mid_block(UpperCAmelCase			,     UpperCAmelCase	)
         lowerCamelCase__    :      Optional[int]					 =		sample.to(UpperCAmelCase	)

         # up
         for up_block in self.up_blocks:
             lowerCamelCase__    :      List[Any]					 =		up_block(UpperCAmelCase			,     UpperCAmelCase	)

        # post-process
     if latent_embeds is None:
         lowerCamelCase__    :      str					 =		self.conv_norm_out(UpperCAmelCase	)
     else:
         lowerCamelCase__    :      Optional[Any]					 =		self.conv_norm_out(UpperCAmelCase			,     UpperCAmelCase	)
     lowerCamelCase__    :      List[str]					 =		self.conv_act(UpperCAmelCase	)
     lowerCamelCase__    :      Optional[Any]					 =		self.conv_out(UpperCAmelCase	)

     return sample


class    lowerCAmelCase							(					nn.Module			):
 def __init__(      self  :      Dict			,     UpperCAmelCase  :      str			,     UpperCAmelCase  :      Optional[int]			,     UpperCAmelCase  :      Optional[Any]			,     UpperCAmelCase  :      str=None			,     UpperCAmelCase  :      Optional[Any]="random"			,     UpperCAmelCase  :      Optional[int]=False			,     UpperCAmelCase  :      List[Any]=True	)			->	Optional[int]:
     super().__init__()
     lowerCamelCase__    :      int					 =		n_e
     lowerCamelCase__    :      Any					 =		vq_embed_dim
     lowerCamelCase__    :      Optional[int]					 =		beta
     lowerCamelCase__    :      Optional[Any]					 =		legacy

     lowerCamelCase__    :      Any					 =		nn.Embedding(self.n_e			,     self.vq_embed_dim	)
     self.embedding.weight.data.uniform_(-1.0 / self.n_e			,     1.0 / self.n_e	)

     lowerCamelCase__    :      Any					 =		remap
     if self.remap is not None:
         self.register_buffer('used'			,     torch.tensor(np.load(self.remap	)	)	)
         lowerCamelCase__    :      str					 =		self.used.shape[0]
         lowerCamelCase__    :      Any					 =		unknown_index  # "random" or "extra" or integer
         if self.unknown_index == "extra":
             lowerCamelCase__    :      List[Any]					 =		self.re_embed
             lowerCamelCase__    :      List[Any]					 =		self.re_embed + 1
         print(
             F"""Remapping {self.n_e} indices to {self.re_embed} indices. """
             F"""Using {self.unknown_index} for unknown indices."""	)
     else:
         lowerCamelCase__    :      List[Any]					 =		n_e

     lowerCamelCase__    :      Optional[Any]					 =		sane_index_shape
 def      A_	(      self  :      Union[str, Any]			,     UpperCAmelCase  :      Any	)			->	Dict:
     lowerCamelCase__    :      Any					 =		inds.shape
     assert len(UpperCAmelCase	) > 1
     lowerCamelCase__    :      str					 =		inds.reshape(ishape[0]			,     -1	)
     lowerCamelCase__    :      str					 =		self.used.to(UpperCAmelCase	)
     lowerCamelCase__    :      Dict					 =		(inds[:, :, None] == used[None, None, ...]).long()
     lowerCamelCase__    :      Optional[int]					 =		match.argmax(-1	)
     lowerCamelCase__    :      Dict					 =		match.sum(2	) < 1
     if self.unknown_index == "random":
         lowerCamelCase__    :      Optional[int]					 =		torch.randint(0			,     self.re_embed			,     size=new[unknown].shape	).to(device=new.device	)
     else:
         lowerCamelCase__    :      int					 =		self.unknown_index
     return new.reshape(UpperCAmelCase	)
 def      A_	(      self  :      List[str]			,     UpperCAmelCase  :      Tuple	)			->	Optional[int]:
     lowerCamelCase__    :      List[Any]					 =		inds.shape
     assert len(UpperCAmelCase	) > 1
     lowerCamelCase__    :      List[Any]					 =		inds.reshape(ishape[0]			,     -1	)
     lowerCamelCase__    :      str					 =		self.used.to(UpperCAmelCase	)
     if self.re_embed > self.used.shape[0]:  # extra token
         lowerCamelCase__    :      Any					 =		0  # simply set to zero
     lowerCamelCase__    :      Any					 =		torch.gather(used[None, :][inds.shape[0] * [0], :]			,     1			,     UpperCAmelCase	)
     return back.reshape(UpperCAmelCase	)
 def      A_	(      self  :      Tuple			,     UpperCAmelCase  :      Union[str, Any]	)			->	Optional[int]:
     # reshape z -> (batch, height, width, channel) and flatten
     lowerCamelCase__    :      Tuple					 =		z.permute(0			,     2			,     3			,     1	).contiguous()
     lowerCamelCase__    :      List[Any]					 =		z.view(-1			,     self.vq_embed_dim	)

     # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
     lowerCamelCase__    :      List[str]					 =		torch.argmin(torch.cdist(UpperCAmelCase			,     self.embedding.weight	)			,     dim=1	)

     lowerCamelCase__    :      Optional[Any]					 =		self.embedding(UpperCAmelCase	).view(z.shape	)
     lowerCamelCase__    :      int					 =		None
     lowerCamelCase__    :      Union[str, Any]					 =		None

     # compute loss for embedding
     if not self.legacy:
         lowerCamelCase__    :      Dict					 =		self.beta * torch.mean((z_q.detach() - z) ** 2	) + torch.mean((z_q - z.detach()) ** 2	)
     else:
         lowerCamelCase__    :      List[str]					 =		torch.mean((z_q.detach() - z) ** 2	) + self.beta * torch.mean((z_q - z.detach()) ** 2	)

     # preserve gradients
     lowerCamelCase__    :      Union[str, Any]					 =		z + (z_q - z).detach()

     # reshape back to match original input shape
     lowerCamelCase__    :      Any					 =		z_q.permute(0			,     3			,     1			,     2	).contiguous()

     if self.remap is not None:
         lowerCamelCase__    :      int					 =		min_encoding_indices.reshape(z.shape[0]			,     -1	)  # add batch axis
         lowerCamelCase__    :      Union[str, Any]					 =		self.remap_to_used(UpperCAmelCase	)
         lowerCamelCase__    :      Optional[int]					 =		min_encoding_indices.reshape(-1			,     1	)  # flatten

     if self.sane_index_shape:
         lowerCamelCase__    :      Any					 =		min_encoding_indices.reshape(z_q.shape[0]			,     z_q.shape[2]			,     z_q.shape[3]	)

     return z_q, loss, (perplexity, min_encodings, min_encoding_indices)







 def      A_	(      self  :      int			,     UpperCAmelCase  :      Dict			,     UpperCAmelCase  :      List[Any]	)			->	List[Any]:
     # shape specifying (batch, height, width, channel)
     if self.remap is not None:
         lowerCamelCase__    :      str					 =		indices.reshape(shape[0]			,     -1	)  # add batch axis
         lowerCamelCase__    :      Optional[Any]					 =		self.unmap_to_all(UpperCAmelCase	)
         lowerCamelCase__    :      Union[str, Any]					 =		indices.reshape(-1	)  # flatten again

     # get quantized latent vectors
     lowerCamelCase__    :      str					 =		self.embedding(UpperCAmelCase	)

     if shape is not None:
         lowerCamelCase__    :      List[str]					 =		z_q.view(UpperCAmelCase	)
         # reshape back to match original input shape
         lowerCamelCase__    :      Union[str, Any]					 =		z_q.permute(0			,     3			,     1			,     2	).contiguous()

     return z_q


class    lowerCAmelCase							(					__UpperCamelCase			):
 def __init__(      self  :      Any			,     UpperCAmelCase  :      Tuple			,     UpperCAmelCase  :      List[Any]=False	)			->	List[str]:
     lowerCamelCase__    :      List[str]					 =		parameters
     lowerCamelCase__ ,							lowerCamelCase__    :      List[str]					 =		torch.chunk(UpperCAmelCase			,     2			,     dim=1	)
     lowerCamelCase__    :      Union[str, Any]					 =		torch.clamp(self.logvar			,     -3_0.0			,     2_0.0	)
     lowerCamelCase__    :      List[Any]					 =		deterministic
     lowerCamelCase__    :      Optional[int]					 =		torch.exp(0.5 * self.logvar	)
     lowerCamelCase__    :      Optional[Any]					 =		torch.exp(self.logvar	)
     if self.deterministic:
         lowerCamelCase__    :      str					 =		torch.zeros_like(
             self.mean			,     device=self.parameters.device			,     dtype=self.parameters.dtype	)
 def      A_	(      self  :      List[str]			,     UpperCAmelCase  :      Optional[torch.Generator] = None	)			->	torch.FloatTensor:
     # make sure sample is on the same device as the parameters and has same dtype
     lowerCamelCase__    :      Any					 =		randn_tensor(
         self.mean.shape			,     generator=UpperCAmelCase			,     device=self.parameters.device			,     dtype=self.parameters.dtype	)
     lowerCamelCase__    :      Optional[int]					 =		self.mean + self.std * sample
     return x
 def      A_	(      self  :      Dict			,     UpperCAmelCase  :      Optional[Any]=None	)			->	Optional[Any]:
     if self.deterministic:
         return torch.Tensor([0.0]	)
     else:
         if other is None:
             return 0.5 * torch.sum(torch.pow(self.mean			,     2	) + self.var - 1.0 - self.logvar			,     dim=[1, 2, 3]	)
         else:
             return 0.5 * torch.sum(
                 torch.pow(self.mean - other.mean			,     2	) / other.var
                 + self.var / other.var
                 - 1.0
                 - self.logvar
                 + other.logvar			,     dim=[1, 2, 3]			,     )
 def      A_	(      self  :      str			,     UpperCAmelCase  :      int			,     UpperCAmelCase  :      List[Any]=[1, 2, 3]	)			->	List[str]:
     if self.deterministic:
         return torch.Tensor([0.0]	)
     lowerCamelCase__    :      int					 =		np.log(2.0 * np.pi	)
     return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean			,     2	) / self.var			,     dim=UpperCAmelCase	)







 def      A_	(      self  :      Optional[int]	)			->	List[str]:
     return self.mean









"""simple docstring"""



# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os

from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES
from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
from ...utils.imports import is_botoa_available
from .config_args import SageMakerConfig
from .config_utils import (
    DYNAMO_BACKENDS,
    _ask_field,
    _ask_options,
    _convert_dynamo_backend,
    _convert_mixed_precision,
    _convert_sagemaker_distributed_mode,
    _convert_yes_no_to_bool,
)


if is_botoa_available():
  import botoa  # noqa: F401







def   __lowerCAmelCase						(						lowercase  :  List[str] )	->						str:


    """simple docstring"""





    snake_case    :      Optional[int]      					=    botoa.client("iam" )

    snake_case    :      Any      					=    {
        "Version": "2012-10-17",
        "Statement": [
            {"Effect": "Allow", "Principal": {"Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"}
        ],
    }
    try:
        # create the role, associated with the chosen trust policy
        iam_client.create_role(
            RoleName=lowercase						,							AssumeRolePolicyDocument=json.dumps(lowercase						,							indent=2 ) )
        snake_case    :      Union[str, Any]      					=    {
            "Version": "2012-10-17",
            "Statement": [
                {
                    "Effect": "Allow",
                    "Action": [
                        "sagemaker:*",
                        "ecr:GetDownloadUrlForLayer",
                        "ecr:BatchGetImage",
                        "ecr:BatchCheckLayerAvailability",
                        "ecr:GetAuthorizationToken",
                        "cloudwatch:PutMetricData",
                        "cloudwatch:GetMetricData",
                        "cloudwatch:GetMetricStatistics",
                        "cloudwatch:ListMetrics",
                        "logs:CreateLogGroup",
                        "logs:CreateLogStream",
                        "logs:DescribeLogStreams",
                        "logs:PutLogEvents",
                        "logs:GetLogEvents",
                        "s3:CreateBucket",
                        "s3:ListBucket",
                        "s3:GetBucketLocation",
                        "s3:GetObject",
                        "s3:PutObject",
                    ],
                    "Resource": "*",
                }
            ],
        }
        # attach policy to role
        iam_client.put_role_policy(
            RoleName=lowercase						,							PolicyName=F'{role_name}_policy_permission'						,							PolicyDocument=json.dumps(lowercase						,							indent=2 )						,							)
    except iam_client.exceptions.EntityAlreadyExistsException:
        print(F'role {role_name} already exists. Using existing one' )







def   __lowerCAmelCase						(						lowercase  :  Dict )	->						Optional[int]:


    """simple docstring"""





    snake_case    :      Any      					=    botoa.client("iam" )
    return iam_client.get_role(RoleName=lowercase )["Role"]["Arn"]







def   __lowerCAmelCase						(						)	->						Union[str, Any]:


    """simple docstring"""





    snake_case    :      Optional[int]      					=    _ask_options(
        "How do you want to authorize?"						,							["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "]						,							lowercase						,							)
    snake_case    :      int      					=    None
    if credentials_configuration == 0:
        snake_case    :      Any      					=    _ask_field("Enter your AWS Profile name: [default] "						,							default="default" )
        snake_case    :      List[str]      					=    aws_profile
    else:
        print(
            "Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with,"
            "`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`" )
        snake_case    :      Any      					=    _ask_field("AWS Access Key ID: " )
        snake_case    :      List[str]      					=    aws_access_key_id

        snake_case    :      Optional[int]      					=    _ask_field("AWS Secret Access Key: " )
        snake_case    :      Union[str, Any]      					=    aws_secret_access_key

    snake_case    :      Optional[Any]      					=    _ask_field("Enter your AWS Region: [us-east-1]"						,							default="us-east-1" )
    snake_case    :      List[str]      					=    aws_region

    snake_case    :      List[str]      					=    _ask_options(
        "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?"						,							["Provide IAM Role name", "Create new IAM role using credentials"]						,							lowercase						,							)
    if role_management == 0:
        snake_case    :      Tuple      					=    _ask_field("Enter your IAM role name: " )
    else:
        snake_case    :      Union[str, Any]      					=    "accelerate_sagemaker_execution_role"
        print(F'Accelerate will create an iam role "{iam_role_name}" using the provided credentials' )
        _create_iam_role_for_sagemaker(lowercase )

    snake_case    :      Union[str, Any]      					=    _ask_field(
        "Do you want to use custom Docker image? [yes/NO]: "						,							_convert_yes_no_to_bool						,							default=lowercase						,							error_message="Please enter yes or no."						,							)
    snake_case    :      Any      					=    None
    if is_custom_docker_image:
        snake_case    :      Union[str, Any]      					=    _ask_field("Enter your Docker image: "						,							lambda lowercase : str(lowercase ).lower() )

    snake_case    :      List[Any]      					=    _ask_field(
        "Do you want to provide SageMaker input channels with data locations? [yes/NO]: "						,							_convert_yes_no_to_bool						,							default=lowercase						,							error_message="Please enter yes or no."						,							)
    snake_case    :      List[str]      					=    None
    if is_sagemaker_inputs_enabled:
        snake_case    :      Dict      					=    _ask_field(
            "Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): "						,							lambda lowercase : str(lowercase ).lower()						,							)

    snake_case    :      Tuple      					=    _ask_field(
        "Do you want to enable SageMaker metrics? [yes/NO]: "						,							_convert_yes_no_to_bool						,							default=lowercase						,							error_message="Please enter yes or no."						,							)
    snake_case    :      int      					=    None
    if is_sagemaker_metrics_enabled:
        snake_case    :      int      					=    _ask_field(
            "Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): "						,							lambda lowercase : str(lowercase ).lower()						,							)

    snake_case    :      str      					=    _ask_options(
        "What is the distributed mode?"						,							["No distributed training", "Data parallelism"]						,							_convert_sagemaker_distributed_mode						,							)
    snake_case    :      Tuple      					=    {}
    snake_case    :      Any      					=    _ask_field(
        "Do you wish to optimize your script with torch dynamo?[yes/NO]:"						,							_convert_yes_no_to_bool						,							default=lowercase						,							error_message="Please enter yes or no."						,							)
    if use_dynamo:
        snake_case    :      Any      					=    "dynamo_"
        snake_case    :      Optional[int]      					=    _ask_options(
            "Which dynamo backend would you like to use?"						,							[x.lower() for x in DYNAMO_BACKENDS]						,							_convert_dynamo_backend						,							default=2						,							)
        snake_case    :      Optional[int]      					=    _ask_field(
            "Do you want to customize the defaults sent to torch.compile? [yes/NO]: "						,							_convert_yes_no_to_bool						,							default=lowercase						,							error_message="Please enter yes or no."						,							)

        if use_custom_options:
            snake_case    :      Dict      					=    _ask_options(
                "Which mode do you want to use?"						,							lowercase						,							lambda lowercase : TORCH_DYNAMO_MODES[int(lowercase )]						,							default="default"						,							)
            snake_case    :      Union[str, Any]      					=    _ask_field(
                "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: "						,							_convert_yes_no_to_bool						,							default=lowercase						,							error_message="Please enter yes or no."						,							)
            snake_case    :      Dict      					=    _ask_field(
                "Do you want to enable dynamic shape tracing? [yes/NO]: "						,							_convert_yes_no_to_bool						,							default=lowercase						,							error_message="Please enter yes or no."						,							)
    snake_case    :      List[str]      					=    "Which EC2 instance type you want to use for your training?"
    if distributed_type != SageMakerDistributedType.NO:
        snake_case    :      str      					=    _ask_options(
            lowercase						,							lowercase						,							lambda lowercase : SAGEMAKER_PARALLEL_EC2_INSTANCES[int(lowercase )] )
    else:
        eca_instance_query += "? [ml.p3.2xlarge]:"
        snake_case    :      Union[str, Any]      					=    _ask_field(lowercase						,							lambda lowercase : str(lowercase ).lower()						,							default="ml.p3.2xlarge" )

    snake_case    :      Any      					=    1
    if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
        snake_case    :      Dict      					=    _ask_field(
            "How many machines do you want use? [1]: "						,							lowercase						,							default=1						,							)

    snake_case    :      Union[str, Any]      					=    _ask_options(
        "Do you wish to use FP16 or BF16 (mixed precision)?"						,							["no", "fp16", "bf16", "fp8"]						,							_convert_mixed_precision						,							)

    if use_dynamo and mixed_precision == "no":
        print(
            "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts." )

    return SageMakerConfig(
        image_uri=lowercase						,							compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER						,							distributed_type=lowercase						,							use_cpu=lowercase						,							dynamo_config=lowercase						,							eca_instance_type=lowercase						,							profile=lowercase						,							region=lowercase						,							iam_role_name=lowercase						,							mixed_precision=lowercase						,							num_machines=lowercase						,							sagemaker_inputs_file=lowercase						,							sagemaker_metrics_file=lowercase						,							)





