Cognomen commited on
Commit
5898ca5
Β·
1 Parent(s): 36dfe3a

change schedulers the special jax way because of course that's how it needs to be done

Browse files
Files changed (1) hide show
  1. app.py +7 -1
app.py CHANGED
@@ -35,7 +35,13 @@ pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
35
  revision="flax",
36
  dtype=jnp.bfloat16,
37
  )
38
- pipe.scheduler = FlaxDPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
 
 
 
 
 
 
39
  #pipe.enable_model_cpu_offload()
40
  #pipe.enable_xformers_memory_efficient_attention()
41
 
 
35
  revision="flax",
36
  dtype=jnp.bfloat16,
37
  )
38
+ scheduler, scheduler_state = FlaxDPMSolverMultistepScheduler.from_pretrained(
39
+ "./models/wd-1-5-b2-flax",
40
+ subfolder="scheduler"
41
+ )
42
+ params["scheduler"] = scheduler_state
43
+
44
+ #scheduler = FlaxDPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
45
  #pipe.enable_model_cpu_offload()
46
  #pipe.enable_xformers_memory_efficient_attention()
47