From 4df63d2d197f26181758b5108f003f225fe84874 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 30 Jan 2023 10:11:30 +0300
Subject: split samplers into one more files for k-diffusion

---
 modules/sd_samplers_compvis.py | 8 ++++++++
 1 file changed, 8 insertions(+)

(limited to 'modules/sd_samplers_compvis.py')

diff --git a/modules/sd_samplers_compvis.py b/modules/sd_samplers_compvis.py
index 3d35ff72..88541193 100644
--- a/modules/sd_samplers_compvis.py
+++ b/modules/sd_samplers_compvis.py
@@ -1,4 +1,6 @@
 import math
+import ldm.models.diffusion.ddim
+import ldm.models.diffusion.plms
 
 import numpy as np
 import torch
@@ -7,6 +9,12 @@ from modules.shared import state
 from modules import sd_samplers_common, prompt_parser, shared
 
 
+samplers_data_compvis = [
+    sd_samplers_common.SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), [], {}),
+    sd_samplers_common.SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), [], {}),
+]
+
+
 class VanillaStableDiffusionSampler:
     def __init__(self, constructor, sd_model):
         self.sampler = constructor(sd_model)
-- 
cgit v1.2.3