

'''simple docstring'''



import argparse

import intel_extension_for_pytorch as ipex
import torch

from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline


_UpperCamelCase						:  Optional[Any]							       =				argparse.ArgumentParser("Stable Diffusion script with intel optimization", add_help=False)
parser.add_argument("--dpm", action="store_true", help="Enable DPMSolver or not")
parser.add_argument("--steps", default=None, type=int, help="Num inference steps")
_UpperCamelCase						:  Optional[int]							       =				parser.parse_args()


_UpperCamelCase						:  Union[str, Any]							       =				"cpu"
_UpperCamelCase						:  Optional[int]							       =				"a lovely <dicoo> in red dress and hat, in the snowly and brightly night, with many brighly buildings"

_UpperCamelCase						:  int							       =				"path-to-your-trained-model"
_UpperCamelCase						:  List[Any]							       =				StableDiffusionPipeline.from_pretrained(model_id)
if args.dpm:
     _UpperCamelCase						:  str							       =				DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
_UpperCamelCase						:  Union[str, Any]							       =				pipe.to(device)

# to channels last
_UpperCamelCase						:  List[Any]							       =				pipe.unet.to(memory_format=torch.channels_last)
_UpperCamelCase						:  Optional[int]							       =				pipe.vae.to(memory_format=torch.channels_last)
_UpperCamelCase						:  Dict							       =				pipe.text_encoder.to(memory_format=torch.channels_last)
if pipe.requires_safety_checker:
     _UpperCamelCase						:  Dict							       =				pipe.safety_checker.to(memory_format=torch.channels_last)

# optimize with ipex
_UpperCamelCase						:  Optional[Any]							       =				torch.randn(2, 4, 64, 64)
_UpperCamelCase						:  str							       =				torch.rand(1) * 999
_UpperCamelCase						:  Union[str, Any]							       =				torch.randn(2, 77, 768)
_UpperCamelCase						:  Tuple							       =				(sample, timestep, encoder_hidden_status)
try:
     _UpperCamelCase						:  Union[str, Any]							       =				ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True, sample_input=input_example)
except Exception:
     _UpperCamelCase						:  int							       =				ipex.optimize(pipe.unet.eval(), dtype=torch.bfloataa, inplace=True)
_UpperCamelCase						:  Any							       =				ipex.optimize(pipe.vae.eval(), dtype=torch.bfloataa, inplace=True)
_UpperCamelCase						:  Optional[int]							       =				ipex.optimize(pipe.text_encoder.eval(), dtype=torch.bfloataa, inplace=True)
if pipe.requires_safety_checker:
     _UpperCamelCase						:  List[str]							       =				ipex.optimize(pipe.safety_checker.eval(), dtype=torch.bfloataa, inplace=True)

# compute
_UpperCamelCase						:  Dict							       =				666
_UpperCamelCase						:  Any							       =				torch.Generator(device).manual_seed(seed)
_UpperCamelCase						:  int							       =				{"generator": generator}
if args.steps is not None:
     _UpperCamelCase						:  Union[str, Any]							       =				args.steps

with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloataa):
     _UpperCamelCase						:  List[str]							       =				pipe(prompt, **generate_kwargs).images[0]

# save image
image.save("generated.png")








'''simple docstring'''



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   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      super().tearDown()
      gc.collect()

  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxControlNetModel.from_pretrained(
          'lllyasviel/sd-controlnet-canny'      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxStableDiffusionControlNetPipeline.from_pretrained(
          'runwayml/stable-diffusion-v1-5'      , controlnet=_SCREAMING_SNAKE_CASE      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase 						=							controlnet_params

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

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

      lowerCAmelCase 						=							jax.random.PRNGKey(0	)
      lowerCAmelCase 						=							jax.random.split(_SCREAMING_SNAKE_CASE      , jax.device_count()	)

      lowerCAmelCase 						=							replicate(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)

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

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

      lowerCAmelCase 						=							jnp.asarray(jax.device_get(image_slice.flatten()	)	)
      lowerCAmelCase 						=							jnp.array(
          [0.167_969, 0.116_699, 0.081_543, 0.154_297, 0.132_812, 0.108_887, 0.169_922, 0.169_922, 0.205_078]	)
      print(F'output_slice: {output_slice}'	)
      assert jnp.abs(output_slice - expected_slice	).max() < 1e-2


  def   _SCREAMING_SNAKE_CASE				(				self	):





      '''simple docstring'''

      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxControlNetModel.from_pretrained(
          'lllyasviel/sd-controlnet-openpose'      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase						,   lowerCAmelCase 						=							FlaxStableDiffusionControlNetPipeline.from_pretrained(
          'runwayml/stable-diffusion-v1-5'      , controlnet=_SCREAMING_SNAKE_CASE      , from_pt=_SCREAMING_SNAKE_CASE      , dtype=jnp.bfloataa	)
      lowerCAmelCase 						=							controlnet_params

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

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

      lowerCAmelCase 						=							jax.random.PRNGKey(0	)
      lowerCAmelCase 						=							jax.random.split(_SCREAMING_SNAKE_CASE      , jax.device_count()	)

      lowerCAmelCase 						=							replicate(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)
      lowerCAmelCase 						=							shard(_SCREAMING_SNAKE_CASE	)

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

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

      lowerCAmelCase 						=							jnp.asarray(jax.device_get(image_slice.flatten()	)	)
      lowerCAmelCase 						=							jnp.array(
          [[0.271_484, 0.261_719, 0.275_391, 0.277_344, 0.279_297, 0.291_016, 0.294_922, 0.302_734, 0.302_734]]	)
      print(F'output_slice: {output_slice}'	)
      assert jnp.abs(output_slice - expected_slice	).max() < 1e-2





