Add lycoris sample code

hunyuan_test.py

@@ -2,23 +2,40 @@
 import torch
 
 from k_diffusion.external import DiscreteVDDPMDenoiser
-from k_diffusion.sampling import sample_euler_ancestral
+from k_diffusion.sampling import sample_euler_ancestral, get_sigmas_exponential
 
 from PIL import Image
 from pytorch_lightning import seed_everything
 
 from library.hunyuan_models import *
 from library.hunyuan_utils import *
 
+from lycoris.kohya import create_network_from_weights
+
 
 PROMPT = """
-Very beautiful Steampunk lady, long silver hair, steampunk outfit and weapon, hyperrealism, photorealistic, 8k, unreal engine
+2girls, multiple girls, 
+
+cesario (umamusume), daring tact (umamusume), umamusume, 
+
+shino (ponjiyuusu), 
+
+aged down, ahoge, animal ears, black hair, blue hair, blue sky, braid, cloud, crown braid, 
+day, green eyes, horse ears, horse girl, horse tail, loafers, long hair, mole, mole under eye, 
+multicolored hair, outdoors, pleated skirt, purple eyes, purple sailor collar, purple shirt, 
+purple skirt, purple thighhighs, sailor collar, school uniform, shirt, shoes, short hair, 
+skirt, sky, tail, thighhighs, time paradox, tracen school uniform, two-tone hair, white hair, 
+winter uniform, 
+
+masterpiece, absurdres, latest, safe
 """
 NEG_PROMPT = "错误的眼睛，糟糕的人脸，毁容，糟糕的艺术，变形，多余的肢体，模糊的颜色，模糊，重复，病态，残缺"
-CLIP_TOKENS = 75 * 1 + 2
+CLIP_TOKENS = 75 * 3 + 2
 ATTN_MODE = "xformers"
+H = 1024
+W = 1024
 STEPS = 16
-CFG_SCALE = 7
+CFG_SCALE = 5
 DEVICE = "cuda"
 DTYPE = torch.float16
 
@@ -43,6 +60,21 @@
         vae.requires_grad_(False)
         mt5_embedder.to(torch.float16)
 
+        lycoris_net, state_dict = create_network_from_weights(
+            multiplier=1.0,
+            file="./output/lokr-hy-test-step00000600.safetensors",
+            vae=vae,
+            text_encoder=[clip_encoder, mt5_embedder],
+            unet=denoiser,
+        )
+        lycoris_net.merge_to(
+            text_encoder=[clip_encoder, mt5_embedder],
+            unet=denoiser,
+            weights_sd=state_dict,
+            dtype=DTYPE,
+            device=DEVICE,
+        )
+
         with torch.autocast("cuda"):
             clip_h, clip_m, mt5_h, mt5_m = get_cond(
                 PROMPT,
@@ -68,9 +100,9 @@
 
         style = torch.as_tensor([0] * 2, device=DEVICE)
         # src hw, dst hw, 0, 0
-        size_cond = [1024, 1024, 1024, 1024, 0, 0]
+        size_cond = [H, W, H, W, 0, 0]
         image_meta_size = torch.as_tensor([size_cond] * 2, device=DEVICE)
-        freqs_cis_img = calc_rope(1024, 1024, patch_size, head_dim)
+        freqs_cis_img = calc_rope(H, W, patch_size, head_dim)
 
         denoiser_wrapper = DiscreteVDDPMDenoiser(
             # A quick patch for learn_sigma
@@ -95,7 +127,10 @@ def cfg_denoise_func(x, sigma):
             return uncond + (cond - uncond) * CFG_SCALE
 
         sigmas = denoiser_wrapper.get_sigmas(STEPS).to(DEVICE)
-        x1 = torch.randn(1, 4, 128, 128, dtype=torch.float16, device=DEVICE)
+        sigmas = get_sigmas_exponential(
+            STEPS, denoiser_wrapper.sigma_min, denoiser_wrapper.sigma_max, DEVICE
+        )
+        x1 = torch.randn(1, 4, H//8, W//8, dtype=torch.float16, device=DEVICE)
 
         with torch.autocast("cuda"):
             sample = sample_euler_ancestral(