




import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, 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 				lowerCAmelCase_       (					unittest.TestCase							):






				def 		snake_case	(			self ):
							# clean up the VRAM after each test
							super().tearDown()
							gc.collect()






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_  ,    SCREAMING_SNAKE_CASE_   : List[Any]					  =      FlaxStableDiffusionPipeline.from_pretrained(
							    'stabilityai/stable-diffusion-2'      ,revision='bf16'      ,dtype=jnp.bfloataa      ,)

							SCREAMING_SNAKE_CASE_   : List[str]					  =      'A painting of a squirrel eating a burger'
							SCREAMING_SNAKE_CASE_   : str					  =      jax.device_count()
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      num_samples * [prompt]
							SCREAMING_SNAKE_CASE_   : Any					  =      sd_pipe.prepare_inputs(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      replicate(snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      shard(snake_case__ )

							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      jax.random.PRNGKey(0 )
							SCREAMING_SNAKE_CASE_   : Dict					  =      jax.random.split(snake_case__      ,jax.device_count() )

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      sd_pipe(snake_case__      ,snake_case__      ,snake_case__      ,num_inference_steps=25      ,jit=snake_case__ )[0]
							assert images.shape == (jax.device_count(), 1, 768, 768, 3)

							SCREAMING_SNAKE_CASE_   : Tuple					  =      images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:] )
							SCREAMING_SNAKE_CASE_   : int					  =      images[0, 253:256, 253:256, -1]

							SCREAMING_SNAKE_CASE_   : List[str]					  =      jnp.asarray(jax.device_get(image_slice.flatten() ) )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512] )
							print(F'output_slice: {output_slice}' )
							assert jnp.abs(output_slice - expected_slice ).max() < 1E-2






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      'stabilityai/stable-diffusion-2'
							SCREAMING_SNAKE_CASE_  ,    SCREAMING_SNAKE_CASE_   : List[Any]					  =      FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case__      ,subfolder='scheduler' )
							SCREAMING_SNAKE_CASE_  ,    SCREAMING_SNAKE_CASE_   : Optional[int]					  =      FlaxStableDiffusionPipeline.from_pretrained(
							    snake_case__      ,scheduler=snake_case__      ,revision='bf16'      ,dtype=jnp.bfloataa      ,)
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      scheduler_params

							SCREAMING_SNAKE_CASE_   : Tuple					  =      'A painting of a squirrel eating a burger'
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      jax.device_count()
							SCREAMING_SNAKE_CASE_   : Tuple					  =      num_samples * [prompt]
							SCREAMING_SNAKE_CASE_   : List[str]					  =      sd_pipe.prepare_inputs(snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      replicate(snake_case__ )
							SCREAMING_SNAKE_CASE_   : Any					  =      shard(snake_case__ )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      jax.random.PRNGKey(0 )
							SCREAMING_SNAKE_CASE_   : str					  =      jax.random.split(snake_case__      ,jax.device_count() )

							SCREAMING_SNAKE_CASE_   : Dict					  =      sd_pipe(snake_case__      ,snake_case__      ,snake_case__      ,num_inference_steps=25      ,jit=snake_case__ )[0]
							assert images.shape == (jax.device_count(), 1, 768, 768, 3)

							SCREAMING_SNAKE_CASE_   : List[str]					  =      images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:] )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      images[0, 253:256, 253:256, -1]

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      jnp.asarray(jax.device_get(image_slice.flatten() ) )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297] )
							print(F'output_slice: {output_slice}' )
							assert jnp.abs(output_slice - expected_slice ).max() < 1E-2










import math


def       UpperCamelCase__(			UpperCamelCase__     :			int    )->list:
       A__		   =		[True] * n
       A__		   =		False
       A__		   =		False
       A__		   =		True

       for i in range(3					,	int(n**0.5 + 1    )					,	2    ):
              A__		   =		i * 2
              while index < n:
                     A__		   =		False
                     A__		   =		index + i

       A__		   =		[2]

       for i in range(3					,	UpperCamelCase__					,	2    ):
              if is_prime[i]:
                     primes.append(UpperCamelCase__    )

       return primes


def       UpperCamelCase__(			UpperCamelCase__     :			int = 99_99_66_66_33_33    )->int:
       A__		   =		math.floor(math.sqrt(UpperCamelCase__    )    ) + 1_00
       A__		   =		prime_sieve(UpperCamelCase__    )

       A__		   =		0
       A__		   =		0
       A__		   =		primes[prime_index]

       while (last_prime**2) <= limit:
              A__		   =		primes[prime_index + 1]

              A__		   =		last_prime**2
              A__		   =		next_prime**2

              # Get numbers divisible by lps(current)
              A__		   =		lower_bound + last_prime
              while upper_bound > current <= limit:
                     matches_sum += current
                     current += last_prime

              # Reset the upper_bound
              while (upper_bound - next_prime) > limit:
                     upper_bound -= next_prime

              # Add the numbers divisible by ups(current)
              A__		   =		upper_bound - next_prime
              while current > lower_bound:
                     matches_sum += current
                     current -= next_prime

              # Remove the numbers divisible by both ups and lps
              A__		   =		0
              while upper_bound > current <= limit:
                     if current <= lower_bound:
                            # Increment the current number
                            current += last_prime * next_prime
                            continue

                     if current > limit:
                            break

                     # Remove twice since it was added by both ups and lps
                     matches_sum -= current * 2

                     # Increment the current number
                     current += last_prime * next_prime

              # Setup for next pair
              A__		   =		next_prime
              prime_index += 1

       return matches_sum


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





