@@ -189,7 +189,7 @@ class ModelSamplingContinuousEDM:
189189 @classmethod
190190 def INPUT_TYPES (s ):
191191 return {"required" : { "model" : ("MODEL" ,),
192- "sampling" : (["v_prediction" , "edm" , "edm_playground_v2.5" , "eps" ],),
192+ "sampling" : (["v_prediction" , "edm" , "edm_playground_v2.5" , "eps" , "cosmos_rflow" ],),
193193 "sigma_max" : ("FLOAT" , {"default" : 120.0 , "min" : 0.0 , "max" : 1000.0 , "step" :0.001 , "round" : False }),
194194 "sigma_min" : ("FLOAT" , {"default" : 0.002 , "min" : 0.0 , "max" : 1000.0 , "step" :0.001 , "round" : False }),
195195 }}
@@ -202,6 +202,7 @@ def INPUT_TYPES(s):
202202 def patch (self , model , sampling , sigma_max , sigma_min ):
203203 m = model .clone ()
204204
205+ sampling_base = comfy .model_sampling .ModelSamplingContinuousEDM
205206 latent_format = None
206207 sigma_data = 1.0
207208 if sampling == "eps" :
@@ -215,8 +216,11 @@ def patch(self, model, sampling, sigma_max, sigma_min):
215216 sampling_type = comfy .model_sampling .EDM
216217 sigma_data = 0.5
217218 latent_format = comfy .latent_formats .SDXL_Playground_2_5 ()
219+ elif sampling == "cosmos_rflow" :
220+ sampling_type = comfy .model_sampling .COSMOS_RFLOW
221+ sampling_base = comfy .model_sampling .ModelSamplingCosmosRFlow
218222
219- class ModelSamplingAdvanced (comfy . model_sampling . ModelSamplingContinuousEDM , sampling_type ):
223+ class ModelSamplingAdvanced (sampling_base , sampling_type ):
220224 pass
221225
222226 model_sampling = ModelSamplingAdvanced (model .model .model_config )
0 commit comments