Thomas Male
commited on
Update handler.py
Browse files- handler.py +2 -1
handler.py
CHANGED
|
@@ -32,7 +32,8 @@ class EndpointHandler():
|
|
| 32 |
self.base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[self.base_name])
|
| 33 |
|
| 34 |
print('creating image model...')
|
| 35 |
-
|
|
|
|
| 36 |
self.base_image_model = model_from_config(MODEL_CONFIGS[self.base_image_name], device)
|
| 37 |
self.base_image_model.eval()
|
| 38 |
self.base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[self.base_image_name])
|
|
|
|
| 32 |
self.base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[self.base_name])
|
| 33 |
|
| 34 |
print('creating image model...')
|
| 35 |
+
# default - base40M. use base300M or base1B for better results
|
| 36 |
+
self.base_image_name = 'base300M'
|
| 37 |
self.base_image_model = model_from_config(MODEL_CONFIGS[self.base_image_name], device)
|
| 38 |
self.base_image_model.eval()
|
| 39 |
self.base_diffusion = diffusion_from_config(DIFFUSION_CONFIGS[self.base_image_name])
|