Spaces:
Sleeping
Sleeping
File size: 1,000 Bytes
f205d49 41b7f97 f205d49 6a64212 f205d49 6a64212 f205d49 6a64212 f205d49 6a64212 f205d49 41b7f97 f205d49 6a64212 f205d49 41b7f97 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
import jax
import numpy as np
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from diffusers import FlaxStableDiffusionPipeline
model_path = "sabman/map-diffuser-v3"
pipeline, _params = FlaxStableDiffusionPipeline.from_pretrained(model_path, dtype=jax.numpy.bfloat16)
# prompt = "create a map with traffic signals, busway and residential buildings, in water color style"
def generate_images(prompt):
prng_seed = jax.random.PRNGKey(-1)
num_inference_steps = 50
num_samples = jax.device_count()
prompt = num_samples * [prompt]
prompt_ids = pipeline.prepare_inputs(prompt)
# shard inputs and rng
params = replicate(_params)
prng_seed = jax.random.split(prng_seed, jax.device_count())
prompt_ids = shard(prompt_ids)
images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images
images = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
return images[0]
|