STYLE

mosaicml · Aug 29, 2023 · acc0195 · acc0195
1 parent a66f97e
commit acc0195
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/diffusion/inference/inference_model.py b/diffusion/inference/inference_model.py
@@ -35,7 +35,10 @@ class StableDiffusionInference():
             Default: ``None``.
     """
 
-    def __init__(self, model_name: str = 'stabilityai/stable-diffusion-2-base', pretrained: bool = False, prediction_type: str = 'epsilon'):
+    def __init__(self,
+                 model_name: str = 'stabilityai/stable-diffusion-2-base',
+                 pretrained: bool = False,
+                 prediction_type: str = 'epsilon'):
         self.device = torch.cuda.current_device()
 
         model = stable_diffusion_2(