

import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to gradient accumulation.
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


__UpperCAmelCase    				=     1_6
__UpperCAmelCase    				=     3_2







def 	_snake_case		(    SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE = 16				)       ->      Optional[int]:

     """simple docstring"""







     _lowerCAmelCase   :					Optional[int]        =    AutoTokenizer.from_pretrained("bert-base-cased"				)
     _lowerCAmelCase   :					Union[str, Any]        =    load_dataset("glue"		,							"mrpc"				)

     def tokenize_function(SCREAMING_SNAKE_CASE				):
          # max_length=None => use the model max length (it's actually the default)
          _lowerCAmelCase   :					Tuple        =    tokenizer(examples["sentence1"]		,							examples["sentence2"]		,							truncation=SCREAMING_SNAKE_CASE		,							max_length=SCREAMING_SNAKE_CASE				)
          return outputs

     # Apply the method we just defined to all the examples in all the splits of the dataset
     # starting with the main process first:
     with accelerator.main_process_first():
          _lowerCAmelCase   :					Optional[Any]        =    datasets.map(
              SCREAMING_SNAKE_CASE		,							batched=SCREAMING_SNAKE_CASE		,							remove_columns=["idx", "sentence1", "sentence2"]		,							)

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     _lowerCAmelCase   :					Optional[Any]        =    tokenized_datasets.rename_column("label"		,							"labels"				)

     def collate_fn(SCREAMING_SNAKE_CASE				):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          _lowerCAmelCase   :					Tuple        =    128 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               _lowerCAmelCase   :					str        =    16
          elif accelerator.mixed_precision != "no":
               _lowerCAmelCase   :					int        =    8
          else:
               _lowerCAmelCase   :					Any        =    None

          return tokenizer.pad(
              SCREAMING_SNAKE_CASE		,							padding="longest"		,							max_length=SCREAMING_SNAKE_CASE		,							pad_to_multiple_of=SCREAMING_SNAKE_CASE		,							return_tensors="pt"		,							)

     # Instantiate dataloaders.
     _lowerCAmelCase   :					str        =    DataLoader(
         tokenized_datasets["train"]		,							shuffle=SCREAMING_SNAKE_CASE		,							collate_fn=SCREAMING_SNAKE_CASE		,							batch_size=SCREAMING_SNAKE_CASE				)
     _lowerCAmelCase   :					List[Any]        =    DataLoader(
         tokenized_datasets["validation"]		,							shuffle=SCREAMING_SNAKE_CASE		,							collate_fn=SCREAMING_SNAKE_CASE		,							batch_size=SCREAMING_SNAKE_CASE				)

     return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('TESTING_MOCKED_DATALOADERS', None) == "1":
  from accelerate.test_utils.training import mocked_dataloaders

  __UpperCAmelCase    				=     mocked_dataloaders  # noqa: F811







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

     """simple docstring"""







     if os.environ.get("TESTING_MOCKED_DATALOADERS"		,							SCREAMING_SNAKE_CASE				) == "1":
          _lowerCAmelCase   :					Optional[Any]        =    2
     # New Code #
     _lowerCAmelCase   :					str        =    int(args.gradient_accumulation_steps				)
     _lowerCAmelCase   :					Dict        =    int(args.local_sgd_steps				)
     # Initialize accelerator
     _lowerCAmelCase   :					Dict        =    Accelerator(
         cpu=args.cpu		,							mixed_precision=args.mixed_precision		,							gradient_accumulation_steps=SCREAMING_SNAKE_CASE				)
     if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
          raise NotImplementedError("LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)"				)
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     _lowerCAmelCase   :					Dict        =    config["lr"]
     _lowerCAmelCase   :					List[str]        =    int(config["num_epochs"]				)
     _lowerCAmelCase   :					List[str]        =    int(config["seed"]				)
     _lowerCAmelCase   :					str        =    int(config["batch_size"]				)

     _lowerCAmelCase   :					List[Any]        =    evaluate.load("glue"		,							"mrpc"				)

     set_seed(SCREAMING_SNAKE_CASE				)
     _lowerCAmelCase   :					Dict        =    get_dataloaders(SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE				)
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     _lowerCAmelCase   :					Optional[Any]        =    AutoModelForSequenceClassification.from_pretrained("bert-base-cased"		,							return_dict=SCREAMING_SNAKE_CASE				)

     # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
     # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
     # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
     _lowerCAmelCase   :					List[Any]        =    model.to(accelerator.device				)

     # Instantiate optimizer
     _lowerCAmelCase   :					List[Any]        =    AdamW(params=model.parameters()		,							lr=SCREAMING_SNAKE_CASE				)

     # Instantiate scheduler
     _lowerCAmelCase   :					str        =    get_linear_schedule_with_warmup(
         optimizer=SCREAMING_SNAKE_CASE		,							num_warmup_steps=100		,							num_training_steps=(len(SCREAMING_SNAKE_CASE				) * num_epochs)		,							)

     # Prepare everything
     # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
     # prepare method.
     _lowerCAmelCase   :					int        =    accelerator.prepare(
         SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE				)

     # Now we train the model
     for epoch in range(SCREAMING_SNAKE_CASE				):
          model.train()
          with LocalSGD(
              accelerator=SCREAMING_SNAKE_CASE		,							model=SCREAMING_SNAKE_CASE		,							local_sgd_steps=SCREAMING_SNAKE_CASE		,							enabled=local_sgd_steps is not None				) as local_sgd:
               for step, batch in enumerate(SCREAMING_SNAKE_CASE				):
                    # We could avoid this line since we set the accelerator with `device_placement=True`.
                    batch.to(accelerator.device				)
                    # New code #
                    # We use the new `accumulate` context manager to perform gradient accumulation
                    # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
                    with accelerator.accumulate(SCREAMING_SNAKE_CASE				):
                         _lowerCAmelCase   :					Tuple        =    model(**SCREAMING_SNAKE_CASE				)
                         _lowerCAmelCase   :					Union[str, Any]        =    output.loss
                         accelerator.backward(SCREAMING_SNAKE_CASE				)
                         optimizer.step()
                         lr_scheduler.step()
                         optimizer.zero_grad()
                         # LocalSGD-specific line
                         local_sgd.step()

          model.eval()
          for step, batch in enumerate(SCREAMING_SNAKE_CASE				):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device				)
               with torch.no_grad():
                    _lowerCAmelCase   :					List[Any]        =    model(**SCREAMING_SNAKE_CASE				)
               _lowerCAmelCase   :					List[Any]        =    outputs.logits.argmax(dim=-1				)
               _lowerCAmelCase   :					Any        =    accelerator.gather_for_metrics((predictions, batch["labels"])				)
               metric.add_batch(
                   predictions=SCREAMING_SNAKE_CASE		,							references=SCREAMING_SNAKE_CASE		,							)

          _lowerCAmelCase   :					Tuple        =    metric.compute()
          # Use accelerator.print to print only on the main process.
          accelerator.print(f'''epoch {epoch}:'''		,							SCREAMING_SNAKE_CASE				)







def 	_snake_case		(    )       ->      Optional[int]:

     """simple docstring"""







     _lowerCAmelCase   :					Optional[int]        =    argparse.ArgumentParser(description="Simple example of training script."				)
     parser.add_argument(
         "--mixed_precision"		,							type=SCREAMING_SNAKE_CASE		,							default=SCREAMING_SNAKE_CASE		,							choices=["no", "fp16", "bf16", "fp8"]		,							help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU."		,							)
     # New Code #
     parser.add_argument(
         "--gradient_accumulation_steps"		,							type=SCREAMING_SNAKE_CASE		,							default=1		,							help="The number of minibatches to be ran before gradients are accumulated."		,							)
     parser.add_argument(
         "--local_sgd_steps"		,							type=SCREAMING_SNAKE_CASE		,							default=8		,							help="Number of local SGD steps or None to disable local SGD"				)
     parser.add_argument("--cpu"		,							action="store_true"		,							help="If passed, will train on the CPU."				)
     _lowerCAmelCase   :					str        =    parser.parse_args()
     _lowerCAmelCase   :					Any        =    {"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(SCREAMING_SNAKE_CASE		,							SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
  main()









import unittest

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

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    TextToVideoSDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import is_xformers_available, load_numpy, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





@skip_mps
class A__    ( A  ,		unittest.TestCase ):



 """simple docstring"""







 _lowercase     :					str      =						TextToVideoSDPipeline
 _lowercase     :					str      =						TEXT_TO_IMAGE_PARAMS
 _lowercase     :					Dict      =						TEXT_TO_IMAGE_BATCH_PARAMS
 # No `output_type`.
 _lowercase     :					Union[str, Any]      =						frozenset(
     [
         '''num_inference_steps''',
         '''generator''',
         '''latents''',
         '''return_dict''',
         '''callback''',
         '''callback_steps''',
     ] )

 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      torch.manual_seed(0       )
      _lowerCAmelCase   :					Union[str, Any]        =    UNetaDConditionModel(
          block_out_channels=(3_2, 6_4, 6_4, 6_4)     ,						layers_per_block=2     ,						sample_size=3_2     ,						in_channels=4     ,						out_channels=4     ,						down_block_types=("CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "CrossAttnDownBlock3D", "DownBlock3D")     ,						up_block_types=("UpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D", "CrossAttnUpBlock3D")     ,						cross_attention_dim=3_2     ,						attention_head_dim=4     ,						)
      _lowerCAmelCase   :					Any        =    DDIMScheduler(
          beta_start=0.00085     ,						beta_end=0.012     ,						beta_schedule="scaled_linear"     ,						clip_sample=A_     ,						set_alpha_to_one=A_     ,						)
      torch.manual_seed(0       )
      _lowerCAmelCase   :					Dict        =    AutoencoderKL(
          block_out_channels=[3_2, 6_4]     ,						in_channels=3     ,						out_channels=3     ,						down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]     ,						up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]     ,						latent_channels=4     ,						sample_size=1_2_8     ,						)
      torch.manual_seed(0       )
      _lowerCAmelCase   :					Union[str, Any]        =    CLIPTextConfig(
          bos_token_id=0     ,						eos_token_id=2     ,						hidden_size=3_2     ,						intermediate_size=3_7     ,						layer_norm_eps=1E-05     ,						num_attention_heads=4     ,						num_hidden_layers=5     ,						pad_token_id=1     ,						vocab_size=1_0_0_0     ,						hidden_act="gelu"     ,						projection_dim=5_1_2     ,						)
      _lowerCAmelCase   :					Any        =    CLIPTextModel(A_       )
      _lowerCAmelCase   :					Union[str, Any]        =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"       )

      _lowerCAmelCase   :					Dict        =    {
          "unet": unet,
          "scheduler": scheduler,
          "vae": vae,
          "text_encoder": text_encoder,
          "tokenizer": tokenizer,
      }
      return components

 def    __magic_name__       (		self :	str     ,						A_ :	Union[str, Any]     ,						A_ :	List[str]=0       ):






      '''simple docstring'''




      if str(A_       ).startswith("mps"       ):
           _lowerCAmelCase   :					Any        =    torch.manual_seed(A_       )
      else:
           _lowerCAmelCase   :					Optional[int]        =    torch.Generator(device=A_       ).manual_seed(A_       )
      _lowerCAmelCase   :					List[str]        =    {
          "prompt": "A painting of a squirrel eating a burger",
          "generator": generator,
          "num_inference_steps": 2,
          "guidance_scale": 6.0,
          "output_type": "pt",
      }
      return inputs

 def    __magic_name__       (		self :	List[str]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    "cpu"  # ensure determinism for the device-dependent torch.Generator
      _lowerCAmelCase   :					Dict        =    self.get_dummy_components()
      _lowerCAmelCase   :					int        =    TextToVideoSDPipeline(**A_       )
      _lowerCAmelCase   :					Dict        =    sd_pipe.to(A_       )
      sd_pipe.set_progress_bar_config(disable=A_       )

      _lowerCAmelCase   :					Optional[Any]        =    self.get_dummy_inputs(A_       )
      _lowerCAmelCase   :					Any        =    "np"
      _lowerCAmelCase   :					List[str]        =    sd_pipe(**A_       ).frames
      _lowerCAmelCase   :					Optional[int]        =    frames[0][-3:, -3:, -1]

      assert frames[0].shape == (6_4, 6_4, 3)
      _lowerCAmelCase   :					str        =    np.array([158.0, 160.0, 153.0, 125.0, 100.0, 121.0, 111.0, 93.0, 113.0]       )

      assert np.abs(image_slice.flatten() - expected_slice       ).max() < 1E-2

 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




      self._test_attention_slicing_forward_pass(test_mean_pixel_difference=A_     ,						expected_max_diff=3E-3       )

 @unittest.skipIf(
     torch_device != "cuda" or not is_xformers_available()     ,						reason="XFormers attention is only available with CUDA and `xformers` installed"     ,						)
 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      self._test_xformers_attention_forwardGenerator_pass(test_mean_pixel_difference=A_     ,						expected_max_diff=1E-2       )

 @unittest.skip(reason="Batching needs to be properly figured out first for this pipeline."       )
 def    __magic_name__       (		self :	List[str]       ):






      '''simple docstring'''




      pass

 @unittest.skip(reason="Batching needs to be properly figured out first for this pipeline."       )
 def    __magic_name__       (		self :	int       ):






      '''simple docstring'''




      pass

 @unittest.skip(reason="`num_images_per_prompt` argument is not supported for this pipeline."       )
 def    __magic_name__       (		self :	str       ):






      '''simple docstring'''




      pass


 def    __magic_name__       (		self :	Optional[Any]       ):






      '''simple docstring'''




      return super().test_progress_bar()





@slow
@skip_mps
class A__    ( unittest.TestCase ):



 """simple docstring"""







 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[Any]        =    load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video.npy"       )

      _lowerCAmelCase   :					Dict        =    TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b"       )
      _lowerCAmelCase   :					Optional[Any]        =    DPMSolverMultistepScheduler.from_config(pipe.scheduler.config       )
      _lowerCAmelCase   :					Optional[Any]        =    pipe.to("cuda"       )

      _lowerCAmelCase   :					Union[str, Any]        =    "Spiderman is surfing"
      _lowerCAmelCase   :					Any        =    torch.Generator(device="cpu"       ).manual_seed(0       )

      _lowerCAmelCase   :					Union[str, Any]        =    pipe(A_     ,						generator=A_     ,						num_inference_steps=2_5     ,						output_type="pt"       ).frames
      _lowerCAmelCase   :					Optional[Any]        =    video_frames.cpu().numpy()

      assert np.abs(expected_video - video       ).mean() < 5E-2


 def    __magic_name__       (		self :	Union[str, Any]       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    load_numpy(
          "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/text_to_video/video_2step.npy"       )

      _lowerCAmelCase   :					Optional[Any]        =    TextToVideoSDPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b"       )
      _lowerCAmelCase   :					Optional[int]        =    pipe.to("cuda"       )

      _lowerCAmelCase   :					List[Any]        =    "Spiderman is surfing"
      _lowerCAmelCase   :					List[str]        =    torch.Generator(device="cpu"       ).manual_seed(0       )

      _lowerCAmelCase   :					List[Any]        =    pipe(A_     ,						generator=A_     ,						num_inference_steps=2     ,						output_type="pt"       ).frames
      _lowerCAmelCase   :					str        =    video_frames.cpu().numpy()

      assert np.abs(expected_video - video       ).mean() < 5E-2






