Update app.py
Browse files
app.py
CHANGED
|
@@ -111,17 +111,18 @@ def load_and_prepare_model(model_id):
|
|
| 111 |
"ford442/RealVisXL_V5.0_BF16": torch.bfloat16,
|
| 112 |
}
|
| 113 |
dtype = model_dtypes.get(model_id, torch.bfloat16) # Default to float32 if not found
|
|
|
|
|
|
|
| 114 |
vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16,safety_checker=None)
|
| 115 |
pipe = StableDiffusionXLPipeline.from_pretrained(
|
| 116 |
model_id,
|
| 117 |
torch_dtype=torch.bfloat16,
|
| 118 |
add_watermarker=False,
|
| 119 |
use_safetensors=True,
|
|
|
|
| 120 |
vae=vae,
|
| 121 |
).to('cuda')
|
| 122 |
-
pipe.unet.set_attn_processor(AttnProcessor2_0())
|
| 123 |
|
| 124 |
-
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear", algorithm_type="sde-dpmsolver++")
|
| 125 |
return pipe
|
| 126 |
|
| 127 |
# Preload and compile both models
|
|
|
|
| 111 |
"ford442/RealVisXL_V5.0_BF16": torch.bfloat16,
|
| 112 |
}
|
| 113 |
dtype = model_dtypes.get(model_id, torch.bfloat16) # Default to float32 if not found
|
| 114 |
+
sched = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config, beta_schedule="scaled_linear", algorithm_type="sde-dpmsolver++")
|
| 115 |
+
|
| 116 |
vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16,safety_checker=None)
|
| 117 |
pipe = StableDiffusionXLPipeline.from_pretrained(
|
| 118 |
model_id,
|
| 119 |
torch_dtype=torch.bfloat16,
|
| 120 |
add_watermarker=False,
|
| 121 |
use_safetensors=True,
|
| 122 |
+
scheduler=sched,
|
| 123 |
vae=vae,
|
| 124 |
).to('cuda')
|
|
|
|
| 125 |
|
|
|
|
| 126 |
return pipe
|
| 127 |
|
| 128 |
# Preload and compile both models
|