float16

app.py

@@ -7,14 +7,14 @@ import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "black-forest-labs/FLUX.1-schnell" #Replace to the model you would like to use
-#enhanceaiteam/Flux-uncensored
 
 if torch.cuda.is_available():
     torch_dtype = torch.float16
 else:
     torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch.float16)
+pipe.unet.load_attn_procs("enhanceaiteam/Flux-uncensored-V2")
 pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max