
import os
import tempfile
import unittest

import numpy as np

from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import require_flax, slow


if is_flax_available():
						import jax
						import jax.numpy as jnp
						from flax.jax_utils import replicate
						from flax.training.common_utils import shard

						from diffusers import FlaxDDIMScheduler, FlaxDiffusionPipeline, FlaxStableDiffusionPipeline





@require_flax
class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



						def        SCREAMING_SNAKE_CASE(			self							:Dict     )  ->str:
										with tempfile.TemporaryDirectory() as tmpdirname:
														# pipeline has Flax weights
														lowercase         =      FlaxDiffusionPipeline.from_pretrained(
														    "hf-internal-testing/tiny-stable-diffusion-pipe"		,    safety_checker=lowerCAmelCase__		,    cache_dir=lowerCAmelCase__     )

														lowercase         =      [t[-1] for t in os.walk(os.path.join(lowerCAmelCase__		,    os.listdir(lowerCAmelCase__     )[0]		,    "snapshots"     )     )]
														lowercase         =      [item for sublist in all_root_files for item in sublist]

														# None of the downloaded files should be a PyTorch file even if we have some here:
														# https://huggingface.co/hf-internal-testing/tiny-stable-diffusion-pipe/blob/main/unet/diffusion_pytorch_model.bin
														assert not any(f.endswith(".bin"     ) for f in files     )





@slow
@require_flax
class   UpperCamelCase_      (      unittest.TestCase		):







						'''simple docstring'''



						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->Tuple:
										lowercase     ,	lowercase         =      FlaxStableDiffusionPipeline.from_pretrained(
										    "hf-internal-testing/tiny-stable-diffusion-pipe"		,    safety_checker=lowerCAmelCase__     )

										lowercase         =      (
										    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
										    " field, close up, split lighting, cinematic"
										)

										lowercase         =      jax.random.PRNGKey(0     )
										lowercase         =      4

										lowercase         =      jax.device_count()
										lowercase         =      num_samples * [prompt]
										lowercase         =      pipeline.prepare_inputs(lowerCAmelCase__     )

										# shard inputs and rng
										lowercase         =      replicate(lowerCAmelCase__     )
										lowercase         =      jax.random.split(lowerCAmelCase__		,    lowerCAmelCase__     )
										lowercase         =      shard(lowerCAmelCase__     )

										lowercase         =      pipeline(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    jit=lowerCAmelCase__     ).images

										assert images.shape == (num_samples, 1, 64, 64, 3)
										if jax.device_count() == 8:
														assert np.abs(np.abs(images[0, 0, :2, :2, -2:]		,    dtype=np.floataa     ).sum() - 4.1_51_47_45     ) < 1E-3
														assert np.abs(np.abs(lowerCAmelCase__		,    dtype=np.floataa     ).sum() - 4_99_47.8_75     ) < 5E-1

										lowercase         =      pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:]     )     )     )
										assert len(lowerCAmelCase__     ) == num_samples





						def        SCREAMING_SNAKE_CASE(			self							:Dict     )  ->Optional[int]:
										lowercase     ,	lowercase         =      FlaxStableDiffusionPipeline.from_pretrained(
										    "CompVis/stable-diffusion-v1-4"		,    revision="flax"		,    safety_checker=lowerCAmelCase__     )

										lowercase         =      (
										    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
										    " field, close up, split lighting, cinematic"
										)

										lowercase         =      jax.random.PRNGKey(0     )
										lowercase         =      50

										lowercase         =      jax.device_count()
										lowercase         =      num_samples * [prompt]
										lowercase         =      pipeline.prepare_inputs(lowerCAmelCase__     )

										# shard inputs and rng
										lowercase         =      replicate(lowerCAmelCase__     )
										lowercase         =      jax.random.split(lowerCAmelCase__		,    lowerCAmelCase__     )
										lowercase         =      shard(lowerCAmelCase__     )

										lowercase         =      pipeline(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    jit=lowerCAmelCase__     ).images

										assert images.shape == (num_samples, 1, 512, 512, 3)
										if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,    dtype=np.floataa     ).sum() - 0.05_65_24_01)     ) < 1E-3
														assert np.abs((np.abs(lowerCAmelCase__		,    dtype=np.floataa     ).sum() - 2_38_38_08.2)     ) < 5E-1





						def        SCREAMING_SNAKE_CASE(			self							:Any     )  ->Union[str, Any]:
										lowercase     ,	lowercase         =      FlaxStableDiffusionPipeline.from_pretrained(
										    "CompVis/stable-diffusion-v1-4"		,    revision="bf16"		,    dtype=jnp.bfloataa		,    safety_checker=lowerCAmelCase__     )

										lowercase         =      (
										    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
										    " field, close up, split lighting, cinematic"
										)

										lowercase         =      jax.random.PRNGKey(0     )
										lowercase         =      50

										lowercase         =      jax.device_count()
										lowercase         =      num_samples * [prompt]
										lowercase         =      pipeline.prepare_inputs(lowerCAmelCase__     )

										# shard inputs and rng
										lowercase         =      replicate(lowerCAmelCase__     )
										lowercase         =      jax.random.split(lowerCAmelCase__		,    lowerCAmelCase__     )
										lowercase         =      shard(lowerCAmelCase__     )

										lowercase         =      pipeline(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    jit=lowerCAmelCase__     ).images

										assert images.shape == (num_samples, 1, 512, 512, 3)
										if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,    dtype=np.floataa     ).sum() - 0.04_00_39_06)     ) < 1E-3
														assert np.abs((np.abs(lowerCAmelCase__		,    dtype=np.floataa     ).sum() - 2_37_35_16.75)     ) < 5E-1





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->int:
										lowercase     ,	lowercase         =      FlaxStableDiffusionPipeline.from_pretrained(
										    "CompVis/stable-diffusion-v1-4"		,    revision="bf16"		,    dtype=jnp.bfloataa     )

										lowercase         =      (
										    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
										    " field, close up, split lighting, cinematic"
										)

										lowercase         =      jax.random.PRNGKey(0     )
										lowercase         =      50

										lowercase         =      jax.device_count()
										lowercase         =      num_samples * [prompt]
										lowercase         =      pipeline.prepare_inputs(lowerCAmelCase__     )

										# shard inputs and rng
										lowercase         =      replicate(lowerCAmelCase__     )
										lowercase         =      jax.random.split(lowerCAmelCase__		,    lowerCAmelCase__     )
										lowercase         =      shard(lowerCAmelCase__     )

										lowercase         =      pipeline(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    jit=lowerCAmelCase__     ).images

										assert images.shape == (num_samples, 1, 512, 512, 3)
										if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,    dtype=np.floataa     ).sum() - 0.04_00_39_06)     ) < 1E-3
														assert np.abs((np.abs(lowerCAmelCase__		,    dtype=np.floataa     ).sum() - 2_37_35_16.75)     ) < 5E-1





						def        SCREAMING_SNAKE_CASE(			self							:Optional[int]     )  ->Optional[int]:
										lowercase         =      FlaxDDIMScheduler(
										    beta_start=0.0_00_85		,    beta_end=0.0_12		,    beta_schedule="scaled_linear"		,    set_alpha_to_one=lowerCAmelCase__		,    steps_offset=1		,    )

										lowercase     ,	lowercase         =      FlaxStableDiffusionPipeline.from_pretrained(
										    "CompVis/stable-diffusion-v1-4"		,    revision="bf16"		,    dtype=jnp.bfloataa		,    scheduler=lowerCAmelCase__		,    safety_checker=lowerCAmelCase__		,    )
										lowercase         =      scheduler.create_state()

										lowercase         =      scheduler_state

										lowercase         =      (
										    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
										    " field, close up, split lighting, cinematic"
										)

										lowercase         =      jax.random.PRNGKey(0     )
										lowercase         =      50

										lowercase         =      jax.device_count()
										lowercase         =      num_samples * [prompt]
										lowercase         =      pipeline.prepare_inputs(lowerCAmelCase__     )

										# shard inputs and rng
										lowercase         =      replicate(lowerCAmelCase__     )
										lowercase         =      jax.random.split(lowerCAmelCase__		,    lowerCAmelCase__     )
										lowercase         =      shard(lowerCAmelCase__     )

										lowercase         =      pipeline(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    jit=lowerCAmelCase__     ).images

										assert images.shape == (num_samples, 1, 512, 512, 3)
										if jax.device_count() == 8:
														assert np.abs((np.abs(images[0, 0, :2, :2, -2:]		,    dtype=np.floataa     ).sum() - 0.0_45_04_39_45)     ) < 1E-3
														assert np.abs((np.abs(lowerCAmelCase__		,    dtype=np.floataa     ).sum() - 2_34_76_93.5)     ) < 5E-1





						def        SCREAMING_SNAKE_CASE(			self							:Optional[Any]     )  ->Any:
										lowercase         =      (
										    "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of"
										    " field, close up, split lighting, cinematic"
										)

										lowercase         =      jax.device_count()
										lowercase         =      num_samples * [prompt]
										lowercase         =      jax.random.split(jax.random.PRNGKey(0     )		,    lowerCAmelCase__     )

										lowercase     ,	lowercase         =      FlaxStableDiffusionPipeline.from_pretrained(
										    "CompVis/stable-diffusion-v1-4"		,    revision="bf16"		,    dtype=jnp.bfloataa		,    safety_checker=lowerCAmelCase__		,    )

										lowercase         =      replicate(lowerCAmelCase__     )
										lowercase         =      pipeline.prepare_inputs(lowerCAmelCase__     )
										lowercase         =      shard(lowerCAmelCase__     )
										lowercase         =      pipeline(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    jit=lowerCAmelCase__     ).images
										assert images.shape == (num_samples, 1, 512, 512, 3)
										lowercase         =      images[2, 0, 256, 10:17, 1]

										# With memory efficient attention
										lowercase     ,	lowercase         =      FlaxStableDiffusionPipeline.from_pretrained(
										    "CompVis/stable-diffusion-v1-4"		,    revision="bf16"		,    dtype=jnp.bfloataa		,    safety_checker=lowerCAmelCase__		,    use_memory_efficient_attention=lowerCAmelCase__		,    )

										lowercase         =      replicate(lowerCAmelCase__     )
										lowercase         =      pipeline.prepare_inputs(lowerCAmelCase__     )
										lowercase         =      shard(lowerCAmelCase__     )
										lowercase         =      pipeline(lowerCAmelCase__		,    lowerCAmelCase__		,    lowerCAmelCase__		,    jit=lowerCAmelCase__     ).images
										assert images_eff.shape == (num_samples, 1, 512, 512, 3)
										lowercase         =      images[2, 0, 256, 10:17, 1]

										# I checked the results visually and they are very similar. However, I saw that the max diff is `1` and the `sum`
										# over the 8 images is exactly `256`, which is very suspicious. Testing a random slice for now.
										assert abs(slice_eff - slice     ).max() < 1E-2


import argparse
import json
import os

import torch
from transformers.file_utils import has_file

from diffusers import UNetaDConditionModel, UNetaDModel


_snake_case      :							Tuple     	=    False
_snake_case      :							Optional[int]     	=    True
_snake_case      :							Any     	=    False


if __name__ == "__main__":
						_snake_case      :							Any     	=    argparse.ArgumentParser()

						parser.add_argument(
						    "--repo_path",
						    default=None,
						    type=str,
						    required=True,
						    help="The config json file corresponding to the architecture.",
						)

						parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the output model.")

						_snake_case      :							Tuple     	=    parser.parse_args()

						_snake_case      :							Optional[Any]     	=    {
						    "image_size": "sample_size",
						    "num_res_blocks": "layers_per_block",
						    "block_channels": "block_out_channels",
						    "down_blocks": "down_block_types",
						    "up_blocks": "up_block_types",
						    "downscale_freq_shift": "freq_shift",
						    "resnet_num_groups": "norm_num_groups",
						    "resnet_act_fn": "act_fn",
						    "resnet_eps": "norm_eps",
						    "num_head_channels": "attention_head_dim",
						}

						_snake_case      :							Tuple     	=    {
						    "time_steps": "time_proj",
						    "mid": "mid_block",
						    "downsample_blocks": "down_blocks",
						    "upsample_blocks": "up_blocks",
						}

						_snake_case      :							Any     	=    "" if has_file(args.repo_path, "config.json") else "unet"

						with open(os.path.join(args.repo_path, subfolder, "config.json"), "r", encoding="utf-8") as reader:
												_snake_case      :							List[Any]     	=    reader.read()
												_snake_case      :							List[Any]     	=    json.loads(text)

						if do_only_config:
												for key in config_parameters_to_change.keys():
																		config.pop(key, None)

						if has_file(args.repo_path, "config.json"):
												_snake_case      :							List[Any]     	=    UNetaDModel(**config)
						else:
												_snake_case      :							int     	=    UNetaDConditionModel if "ldm-text2im-large-256" in args.repo_path else UNetaDModel
												_snake_case      :							Optional[Any]     	=    class_name(**config)

						if do_only_config:
												model.save_config(os.path.join(args.repo_path, subfolder))

						_snake_case      :							Any     	=    dict(model.config)

						if do_only_renaming:
												for key, value in config_parameters_to_change.items():
																		if key in config:
																								_snake_case      :							Any     	=    config[key]
																								del config[key]

												_snake_case      :							List[str]     	=    [k.replace("UNetRes", "") for k in config["down_block_types"]]
												_snake_case      :							Any     	=    [k.replace("UNetRes", "") for k in config["up_block_types"]]

						if do_only_weights:
												_snake_case      :							Optional[int]     	=    torch.load(os.path.join(args.repo_path, subfolder, "diffusion_pytorch_model.bin"))

												_snake_case      :							Any     	=    {}
												for param_key, param_value in state_dict.items():
																		if param_key.endswith(".op.bias") or param_key.endswith(".op.weight"):
																								continue
																		_snake_case      :							List[Any]     	=    False
																		for key, new_key in key_parameters_to_change.items():
																								if not has_changed and param_key.split(".")[0] == key:
																														_snake_case      :							Dict     	=    param_value
																														_snake_case      :							Any     	=    True
																		if not has_changed:
																								_snake_case      :							int     	=    param_value

												model.load_state_dict(new_state_dict)
												model.save_pretrained(os.path.join(args.repo_path, subfolder))
