Update app.py
Browse files
app.py
CHANGED
@@ -178,9 +178,9 @@ class DiffusionBuilder:
|
|
178 |
def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None, model_type: str = "StableDiffusion", download: bool = True):
|
179 |
with st.spinner(f"{'Downloading' if download else 'Loading'} {model_path}... ⏳"):
|
180 |
if model_type == "StableDiffusion":
|
181 |
-
self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32,
|
182 |
elif model_type == "DDPM":
|
183 |
-
self.pipeline = DDPMPipeline.from_pretrained(model_path, torch_dtype=torch.float32,
|
184 |
self.pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipeline.scheduler.config)
|
185 |
if config:
|
186 |
self.config = config
|
|
|
178 |
def load_model(self, model_path: str, config: Optional[DiffusionConfig] = None, model_type: str = "StableDiffusion", download: bool = True):
|
179 |
with st.spinner(f"{'Downloading' if download else 'Loading'} {model_path}... ⏳"):
|
180 |
if model_type == "StableDiffusion":
|
181 |
+
self.pipeline = StableDiffusionPipeline.from_pretrained(model_path, torch_dtype=torch.float32, local_files_only=not download).to("cpu")
|
182 |
elif model_type == "DDPM":
|
183 |
+
self.pipeline = DDPMPipeline.from_pretrained(model_path, torch_dtype=torch.float32, local_files_only=not download).to("cpu")
|
184 |
self.pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(self.pipeline.scheduler.config)
|
185 |
if config:
|
186 |
self.config = config
|