CompVis
/

stable-diffusion-v1-4

pcuenq HF Staff commited on Dec 6, 2022

Commit

bd73f2a

1 Parent(s): 114c79c

Do not assume 8 devices in JAX (#154)

- Do not assume 8 devices in JAX (e124bbdca2dab1af0cdce19d575f8043eab9341e)

Co-authored-by: Pedro Cuenca <[email protected]>

Files changed (1) hide show

README.md CHANGED Viewed

@@ -154,7 +154,7 @@ prompt_ids = pipeline.prepare_inputs(prompt)
 # shard inputs and rng
 params = replicate(params)
-prng_seed = jax.random.split(prng_seed, 8)
 prompt_ids = shard(prompt_ids)
 images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images
@@ -187,7 +187,7 @@ prompt_ids = pipeline.prepare_inputs(prompt)
 # shard inputs and rng
 params = replicate(params)
-prng_seed = jax.random.split(prng_seed, 8)
 prompt_ids = shard(prompt_ids)
 images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images

 # shard inputs and rng
 params = replicate(params)
+prng_seed = jax.random.split(prng_seed, num_samples)
 prompt_ids = shard(prompt_ids)
 images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images
 # shard inputs and rng
 params = replicate(params)
+prng_seed = jax.random.split(prng_seed, num_samples)
 prompt_ids = shard(prompt_ids)
 images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images