Spaces:

p1atdev
/

danbooru-tags-transformer-v2

Running on Zero

p1atdev commited on May 5

Commit

6ba5764

•

1 Parent(s): 2c5c599

chore: use bf16

Files changed (1) hide show

diffusion.py CHANGED Viewed

@@ -80,7 +80,7 @@ class ImageGenerator:
     def __init__(self, model_name: str = "cagliostrolab/animagine-xl-3.1"):
         self.pipe = StableDiffusionXLPipeline.from_pretrained(
             model_name,
-            torch_dtype=torch.float16,
             use_safetensors=True,
             add_watermarker=False,
             custom_pipeline="lpw_stable_diffusion_xl",

     def __init__(self, model_name: str = "cagliostrolab/animagine-xl-3.1"):
         self.pipe = StableDiffusionXLPipeline.from_pretrained(
             model_name,
+            torch_dtype=torch.bfloat16,
             use_safetensors=True,
             add_watermarker=False,
             custom_pipeline="lpw_stable_diffusion_xl",