Update app.py
Browse files
app.py
CHANGED
@@ -46,9 +46,9 @@ def infer(prompts, negative_prompts, image):
|
|
46 |
|
47 |
params["controlnet"] = controlnet_params
|
48 |
|
49 |
-
num_samples = jax.device_count("
|
50 |
rng = create_key(0)
|
51 |
-
rng = jax.random.split(rng, jax.device_count("
|
52 |
im = canny_filter(image)
|
53 |
canny_image = Image.fromarray(im)
|
54 |
|
|
|
46 |
|
47 |
params["controlnet"] = controlnet_params
|
48 |
|
49 |
+
num_samples = jax.device_count("gpu")
|
50 |
rng = create_key(0)
|
51 |
+
rng = jax.random.split(rng, jax.device_count("gpu"))
|
52 |
im = canny_filter(image)
|
53 |
canny_image = Image.fromarray(im)
|
54 |
|