Cast safety_checker to float32 (#276)

ZachNagengast · web-flow · commit 7fe1bf23b91f · 2023-10-06T20:54:39.000-07:00
diff --git a/python_coreml_stable_diffusion/torch2coreml.py b/python_coreml_stable_diffusion/torch2coreml.py
@@ -926,6 +926,8 @@ def convert_safety_checker(pipe, args):
         )
         return
 
+    pipe.safety_checker = pipe.safety_checker.to(torch.float32)
+    
     im_h = pipe.vae.config.sample_size
     im_w = pipe.vae.config.sample_size
 

Original file line number	Diff line number	Diff line change
`@@ -926,6 +926,8 @@ def convert_safety_checker(pipe, args):`
`926`	`926`	`)`
`927`	`927`	`return`
`928`	`928`
	`929`	`+ pipe.safety_checker = pipe.safety_checker.to(torch.float32)`
	`930`	`+`
`929`	`931`	`im_h = pipe.vae.config.sample_size`
`930`	`932`	`im_w = pipe.vae.config.sample_size`
`931`	`933`