fix type hint

Author: arrjon

bayesflow/networks/diffusion_model/schedules/edm_noise_schedule.py

@@ -1,4 +1,5 @@
 import math
+from typing import Literal
 
 from keras import ops
 
@@ -19,7 +20,11 @@ class EDMNoiseSchedule(NoiseSchedule):
     """
 
     def __init__(
-        self, sigma_data: float = 1.0, sigma_min: float = 1e-4, sigma_max: float = 80.0, variance_type="preserving"
+        self,
+        sigma_data: float = 1.0,
+        sigma_min: float = 1e-4,
+        sigma_max: float = 80.0,
+        variance_type: Literal["preserving", "exploding"] = "preserving",
     ):
         """
         Initialize the EDM noise schedule.
@@ -33,9 +38,8 @@ def __init__(
             The minimum noise level. Only relevant for sampling. Default is 1e-4.
         sigma_max : float, optional
             The maximum noise level. Only relevant for sampling. Default is 80.0.
-        variance_type : str, optional
-            The type of variance to use. One of "preserving", or "exploding". Default is "preserving". Original EDM
-            paper uses "exploding".
+        variance_type : Literal["preserving", "exploding"], optional
+            The type of variance to use. Default is "preserving". Original EDM paper uses "exploding".
         """
         super().__init__(name="edm_noise_schedule", variance_type=variance_type)
         self.sigma_data = sigma_data