File size: 1,106 Bytes
f205d49
 
 
 
ae143ff
f205d49
41b7f97
e01847c
 
 
 
f205d49
6a64212
f205d49
6a64212
f205d49
e01847c
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
30
31
32
import jax
import numpy as np
from flax.jax_utils import replicate
from flax.training.common_utils import shard
from diffusers import DiffusionPipeline

model_path = "sabman/map-diffuser-v3"
# pipeline, _params = FlaxStableDiffusionPipeline.from_pretrained(model_path, dtype=jax.numpy.bfloat16)
pipeline = DiffusionPipeline.from_pretrained(
    model_path,
    from_flax=True, safety_checker=None).to("cuda")

# 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 = 20

    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]