From 18fb3f56749fa87ba1c6486f1289c29b2dc4ebdd Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sat, 6 Apr 2024 12:05:04 +0200
Subject: [PATCH] Quick fix for new comfy version

---
 PixArt/sampler.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/PixArt/sampler.py b/PixArt/sampler.py
index fa2f39c..229a40f 100644
--- a/PixArt/sampler.py
+++ b/PixArt/sampler.py
@@ -2,7 +2,7 @@ import torch
 from .sampling import gaussian_diffusion as gd
 from .sampling.dpm_solver import model_wrapper, DPM_Solver, NoiseScheduleVP
 
-from comfy.sample import prepare_sampling, prepare_noise, cleanup_additional_models, get_models_from_cond
+from comfy.sample import prepare_sampling, prepare_noise
 import comfy.utils
 import latent_preview
 
@@ -69,6 +69,4 @@ def sample_pixart(model, seed, steps, cfg, noise_schedule, noise_schedule_vp, po
 		previewer=previewer,
 	)
 
-	cleanup_additional_models(models)
-	cleanup_additional_models(set(get_models_from_cond(positive, "control")))
 	return samples.detach().cpu().float() * (1 / model.model.latent_format.scale_factor)