Spaces:
Running
Running
Update diffusion_webui/diffusion_models/img2img_app.py
Browse files
diffusion_webui/diffusion_models/img2img_app.py
CHANGED
@@ -15,14 +15,15 @@ class StableDiffusionImage2ImageGenerator:
|
|
15 |
self.pipe = None
|
16 |
|
17 |
def load_model(self, model_path, scheduler):
|
18 |
-
if self.pipe is None:
|
19 |
self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
|
20 |
model_path, safety_checker=None, torch_dtype=torch.float16
|
21 |
)
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
|
|
26 |
|
27 |
return self.pipe
|
28 |
|
|
|
15 |
self.pipe = None
|
16 |
|
17 |
def load_model(self, model_path, scheduler):
|
18 |
+
if self.pipe is None or self.pipe.model_name != stable_model_path or self.pipe.scheduler_name != scheduler:
|
19 |
self.pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
|
20 |
model_path, safety_checker=None, torch_dtype=torch.float16
|
21 |
)
|
22 |
+
|
23 |
+
self.pipe.model_name = stable_model_path
|
24 |
+
self.pipe = get_scheduler(pipe=self.pipe, scheduler=scheduler)
|
25 |
+
self.pipe.to("cuda")
|
26 |
+
self.pipe.enable_xformers_memory_efficient_attention()
|
27 |
|
28 |
return self.pipe
|
29 |
|