Spaces:
Paused
Paused
Update diffusion_webui/diffusion_models/text2img_app.py
Browse files
diffusion_webui/diffusion_models/text2img_app.py
CHANGED
|
@@ -15,10 +15,10 @@ class StableDiffusionText2ImageGenerator:
|
|
| 15 |
|
| 16 |
def load_model(
|
| 17 |
self,
|
| 18 |
-
|
| 19 |
-
|
| 20 |
):
|
| 21 |
-
if self.pipe is None:
|
| 22 |
self.pipe = StableDiffusionPipeline.from_pretrained(
|
| 23 |
model_path, safety_checker=None, torch_dtype=torch.float16
|
| 24 |
)
|
|
@@ -26,6 +26,7 @@ class StableDiffusionText2ImageGenerator:
|
|
| 26 |
self.pipe = get_scheduler(pipe=self.pipe, scheduler=scheduler)
|
| 27 |
self.pipe.to("cuda")
|
| 28 |
self.pipe.enable_xformers_memory_efficient_attention()
|
|
|
|
| 29 |
|
| 30 |
return self.pipe
|
| 31 |
|
|
|
|
| 15 |
|
| 16 |
def load_model(
|
| 17 |
self,
|
| 18 |
+
stable_model_path,
|
| 19 |
+
scheduler_name,
|
| 20 |
):
|
| 21 |
+
if self.pipe is None or self.pipe.model_name != stable_model_path or self.pipe.scheduler_name != scheduler:
|
| 22 |
self.pipe = StableDiffusionPipeline.from_pretrained(
|
| 23 |
model_path, safety_checker=None, torch_dtype=torch.float16
|
| 24 |
)
|
|
|
|
| 26 |
self.pipe = get_scheduler(pipe=self.pipe, scheduler=scheduler)
|
| 27 |
self.pipe.to("cuda")
|
| 28 |
self.pipe.enable_xformers_memory_efficient_attention()
|
| 29 |
+
self.pipe.model_name = stable_model_path
|
| 30 |
|
| 31 |
return self.pipe
|
| 32 |
|