comi-logic/Wan2GP
7
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

LTXV and Flux updates

Browse Source
This commit is contained in:
deepbeepmeep
2025-07-21 09:37:52 +02:00
parent a356c6af4b
commit a1c228054c
36 changed files with 1285 additions and 2288 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
flux/sampling.py
View File

@@ -224,42 +224,48 @@ def prepare_kontext(
if bs == 1 and not isinstance(prompt, str):
bs = len(prompt)
width, height = img_cond.size
aspect_ratio = width / height
if img_cond != None:
width, height = img_cond.size
aspect_ratio = width / height
# Kontext is trained on specific resolutions, using one of them is recommended
_, width, height = min((abs(aspect_ratio - w / h), w, h) for w, h in PREFERED_KONTEXT_RESOLUTIONS)
# Kontext is trained on specific resolutions, using one of them is recommended
_, width, height = min((abs(aspect_ratio - w / h), w, h) for w, h in PREFERED_KONTEXT_RESOLUTIONS)
width = 2 * int(width / 16)
height = 2 * int(height / 16)
width = 2 * int(width / 16)
height = 2 * int(height / 16)
img_cond = img_cond.resize((8 * width, 8 * height), Image.Resampling.LANCZOS)
img_cond = np.array(img_cond)
img_cond = torch.from_numpy(img_cond).float() / 127.5 - 1.0
img_cond = rearrange(img_cond, "h w c -> 1 c h w")
img_cond_orig = img_cond.clone()
img_cond = img_cond.resize((8 * width, 8 * height), Image.Resampling.LANCZOS)
img_cond = np.array(img_cond)
img_cond = torch.from_numpy(img_cond).float() / 127.5 - 1.0
img_cond = rearrange(img_cond, "h w c -> 1 c h w")
img_cond_orig = img_cond.clone()
with torch.no_grad():
img_cond = ae.encode(img_cond.to(device))
with torch.no_grad():
img_cond = ae.encode(img_cond.to(device))
img_cond = img_cond.to(torch.bfloat16)
img_cond = rearrange(img_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
if img_cond.shape[0] == 1 and bs > 1:
img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
img_cond = img_cond.to(torch.bfloat16)
img_cond = rearrange(img_cond, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
if img_cond.shape[0] == 1 and bs > 1:
img_cond = repeat(img_cond, "1 ... -> bs ...", bs=bs)
# image ids are the same as base image with the first dimension set to 1
# instead of 0
img_cond_ids = torch.zeros(height // 2, width // 2, 3)
img_cond_ids[..., 0] = 1
img_cond_ids[..., 1] = img_cond_ids[..., 1] + torch.arange(height // 2)[:, None]
img_cond_ids[..., 2] = img_cond_ids[..., 2] + torch.arange(width // 2)[None, :]
img_cond_ids = repeat(img_cond_ids, "h w c -> b (h w) c", b=bs)
if target_width is None:
target_width = 8 * width
if target_height is None:
target_height = 8 * height
# image ids are the same as base image with the first dimension set to 1
# instead of 0
img_cond_ids = torch.zeros(height // 2, width // 2, 3)
img_cond_ids[..., 0] = 1
img_cond_ids[..., 1] = img_cond_ids[..., 1] + torch.arange(height // 2)[:, None]
img_cond_ids[..., 2] = img_cond_ids[..., 2] + torch.arange(width // 2)[None, :]
img_cond_ids = repeat(img_cond_ids, "h w c -> b (h w) c", b=bs)
if target_width is None:
target_width = 8 * width
if target_height is None:
target_height = 8 * height
img_cond_ids = img_cond_ids.to(device)
else:
img_cond = None
img_cond_ids = None
img_cond_orig = None
img = get_noise(
bs,
target_height,
@@ -271,7 +277,7 @@ def prepare_kontext(
return_dict = prepare(t5, clip, img, prompt)
return_dict["img_cond_seq"] = img_cond
return_dict["img_cond_seq_ids"] = img_cond_ids.to(device)
return_dict["img_cond_seq_ids"] = img_cond_ids
return_dict["img_cond_orig"] = img_cond_orig
return return_dict, target_height, target_width