i

webui.py

@@ -2,49 +2,32 @@ import os
 import random
 import torch
 import numpy as np
+import modules.core as core
 
-from comfy.sd import load_checkpoint_guess_config
-from nodes import VAEDecode, KSamplerAdvanced, EmptyLatentImage, SaveImage, CLIPTextEncode
 from modules.path import modelfile_path
 
 
 xl_base_filename = os.path.join(modelfile_path, 'sd_xl_base_1.0.safetensors')
 xl_refiner_filename = os.path.join(modelfile_path, 'sd_xl_refiner_1.0.safetensors')
 
-xl_base, xl_base_clip, xl_base_vae, xl_base_clipvision = load_checkpoint_guess_config(xl_base_filename)
-del xl_base_clipvision
+xl_base = core.load_model(xl_base_filename)
 
-opCLIPTextEncode = CLIPTextEncode()
-opEmptyLatentImage = EmptyLatentImage()
-opKSamplerAdvanced = KSamplerAdvanced()
-opVAEDecode = VAEDecode()
+positive_conditions = core.encode_prompt_condition(clip=xl_base.clip, prompt='a handsome man in forest')
+negative_conditions = core.encode_prompt_condition(clip=xl_base.clip, prompt='bad, ugly')
 
-with torch.no_grad():
-    positive_conditions = opCLIPTextEncode.encode(clip=xl_base_clip, text='a handsome man in forest')[0]
-    negative_conditions = opCLIPTextEncode.encode(clip=xl_base_clip, text='bad, ugly')[0]
+empty_latent = core.generate_empty_latent(width=1024, height=1024, batch_size=1)
 
-    initial_latent_image = opEmptyLatentImage.generate(width=1024, height=1024, batch_size=1)[0]
+sampled_latent = core.ksample(
+    unet=xl_base.unet,
+    positive_condition=positive_conditions,
+    negative_condition=negative_conditions,
+    latent_image=empty_latent
+)
 
-    samples = opKSamplerAdvanced.sample(
-        add_noise="enable",
-        noise_seed=random.randint(1, 2 ** 64),
-        steps=25,
-        cfg=9,
-        sampler_name="euler",
-        scheduler="normal",
-        start_at_step=0,
-        end_at_step=25,
-        return_with_leftover_noise="enable",
-        model=xl_base,
-        positive=positive_conditions,
-        negative=negative_conditions,
-        latent_image=initial_latent_image,
-    )[0]
+decoded_latent = core.decode_vae(vae=xl_base.vae, latent_image=sampled_latent)
 
-    vae_decoded = opVAEDecode.decode(samples=samples, vae=xl_base_vae)[0]
+images = core.image_to_numpy(decoded_latent)
 
-    for image in vae_decoded:
-        i = 255. * image.cpu().numpy()
-        img = np.clip(i, 0, 255).astype(np.uint8)
-        import cv2
-        cv2.imwrite('a.png', img[:, :, ::-1])
+for image in images:
+    import cv2
+    cv2.imwrite('a.png', image[:, :, ::-1])