fusing/stable-diffusion-flax-new

#!/usr/bin/env python3
from diffusers import FlaxStableDiffusionPipeline
from jax import pmap
import numpy as np
import jax
from flax.jax_utils import replicate
from flax.training.common_utils import shard


prng_seed = jax.random.PRNGKey(0)
num_inference_steps = 50

pipeline, params = FlaxStableDiffusionPipeline.from_pretrained("fusing/stable-diffusion-flax-new", use_auth_token=True)
del params["safety_checker"]

# pmap
p_sample = pmap(pipeline.__call__, static_broadcasted_argnums=(3,))

# prep prompts
prompt = "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic"
num_samples = jax.device_count()
prompt = num_samples * [prompt]
prompt_ids = pipeline.prepare_inputs(prompt)

# replicate
params = replicate(params)
prng_seed = jax.random.split(prng_seed, 8)
prompt_ids = shard(prompt_ids)

# run
images = p_sample(prompt_ids, params, prng_seed, num_inference_steps).images

# get pil images
images_pil = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))

import ipdb; ipdb.set_trace()
print("Images should be good")
# images_pil[0].save(...)