Cognomen commited on
Commit
5407934
Β·
1 Parent(s): 57596b6

use SDE solver

Browse files
Files changed (1) hide show
  1. app.py +2 -1
app.py CHANGED
@@ -1,7 +1,7 @@
1
  import gradio as gr
2
  import jax.numpy as jnp
3
  from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel
4
- from diffusers import UniPCMultistepScheduler
5
  import torch
6
  torch.backends.cuda.matmul.allow_tf32 = True
7
  import torchvision
@@ -34,6 +34,7 @@ pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
34
  controlnet=cnet,
35
  revision="flax",
36
  dtype=jnp.bfloat16,
 
37
  )
38
  #pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
39
  #pipe.enable_model_cpu_offload()
 
1
  import gradio as gr
2
  import jax.numpy as jnp
3
  from diffusers import FlaxStableDiffusionControlNetPipeline, FlaxControlNetModel
4
+ from diffusers import FlaxScoreSdeVeScheduler
5
  import torch
6
  torch.backends.cuda.matmul.allow_tf32 = True
7
  import torchvision
 
34
  controlnet=cnet,
35
  revision="flax",
36
  dtype=jnp.bfloat16,
37
+ scheduler=FlaxScoreSdeVeScheduler()
38
  )
39
  #pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
40
  #pipe.enable_model_cpu_offload()