


def 							__lowercase       (			__lowerCAmelCase							:    int = 1_0_0_0       ):
							return sum(2 * a * ((a - 1) // 2) for a in range(3 ,   n + 1       )       )


if __name__ == "__main__":
						print(solution())








import gc
import unittest

from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
from diffusers.utils import is_flax_available, load_image, slow
from diffusers.utils.testing_utils import require_flax


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







@slow
@require_flax
class 	snake_case_       (unittest.TestCase				):
					def 			lowerCamelCase__(				self       :List[Any]				) ->   str:
												# clean up the VRAM after each test
												super().tearDown()
												gc.collect()
					def 			lowerCamelCase__(				self       :int				) ->   Optional[Any]:
												a__				, a__					=					FlaxControlNetModel.from_pretrained(
												    'lllyasviel/sd-controlnet-canny'      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__				, a__					=					FlaxStableDiffusionControlNetPipeline.from_pretrained(
												    'runwayml/stable-diffusion-v1-5'      ,controlnet=__snake_case      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__					=					controlnet_params

												a__					=					'bird'
												a__					=					jax.device_count()
												a__					=					pipe.prepare_text_inputs([prompts] * num_samples				)

												a__					=					load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'				)
												a__					=					pipe.prepare_image_inputs([canny_image] * num_samples				)

												a__					=					jax.random.PRNGKey(0				)
												a__					=					jax.random.split(__snake_case      ,jax.device_count()				)

												a__					=					replicate(__snake_case				)
												a__					=					shard(__snake_case				)
												a__					=					shard(__snake_case				)

												a__					=					pipe(
												    prompt_ids=__snake_case      ,image=__snake_case      ,params=__snake_case      ,prng_seed=__snake_case      ,num_inference_steps=50      ,jit=__snake_case      ,).images
												assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

												a__					=					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
												a__					=					images[0, 2_53:2_56, 2_53:2_56, -1]

												a__					=					jnp.asarray(jax.device_get(image_slice.flatten()				)				)
												a__					=					jnp.array(
												    [0.16_79_69, 0.11_66_99, 0.08_15_43, 0.15_42_97, 0.13_28_12, 0.10_88_87, 0.16_99_22, 0.16_99_22, 0.20_50_78]				)
												print(F'output_slice: {output_slice}'				)
												assert jnp.abs(output_slice - expected_slice				).max() < 1E-2



					def 			lowerCamelCase__(				self       :Optional[Any]				) ->   Optional[Any]:
												a__				, a__					=					FlaxControlNetModel.from_pretrained(
												    'lllyasviel/sd-controlnet-openpose'      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__				, a__					=					FlaxStableDiffusionControlNetPipeline.from_pretrained(
												    'runwayml/stable-diffusion-v1-5'      ,controlnet=__snake_case      ,from_pt=__snake_case      ,dtype=jnp.bfloataa				)
												a__					=					controlnet_params

												a__					=					'Chef in the kitchen'
												a__					=					jax.device_count()
												a__					=					pipe.prepare_text_inputs([prompts] * num_samples				)

												a__					=					load_image(
												    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/pose.png'				)
												a__					=					pipe.prepare_image_inputs([pose_image] * num_samples				)

												a__					=					jax.random.PRNGKey(0				)
												a__					=					jax.random.split(__snake_case      ,jax.device_count()				)

												a__					=					replicate(__snake_case				)
												a__					=					shard(__snake_case				)
												a__					=					shard(__snake_case				)

												a__					=					pipe(
												    prompt_ids=__snake_case      ,image=__snake_case      ,params=__snake_case      ,prng_seed=__snake_case      ,num_inference_steps=50      ,jit=__snake_case      ,).images
												assert images.shape == (jax.device_count(), 1, 7_68, 5_12, 3)

												a__					=					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]				)
												a__					=					images[0, 2_53:2_56, 2_53:2_56, -1]

												a__					=					jnp.asarray(jax.device_get(image_slice.flatten()				)				)
												a__					=					jnp.array(
												    [[0.27_14_84, 0.26_17_19, 0.27_53_91, 0.27_73_44, 0.27_92_97, 0.29_10_16, 0.29_49_22, 0.30_27_34, 0.30_27_34]]				)
												print(F'output_slice: {output_slice}'				)
												assert jnp.abs(output_slice - expected_slice				).max() < 1E-2




