| import torch
|
|
|
|
|
|
|
|
|
| class SwarmReferenceOnly:
|
| @classmethod
|
| def INPUT_TYPES(s):
|
| return {
|
| "required": {
|
| "model": ("MODEL",),
|
| "reference": ("LATENT",),
|
| "latent": ("LATENT",)
|
| }
|
| }
|
|
|
| CATEGORY = "SwarmUI/sampling"
|
| RETURN_TYPES = ("MODEL", "LATENT")
|
| FUNCTION = "reference_only"
|
| DESCRIPTION = "Applies 'reference only' image-prompting to the generation. Must forward the new model, and the new latent, to work properly."
|
|
|
| def reference_only(self, model, reference, latent):
|
| model_reference = model.clone()
|
| reference["samples"] = torch.nn.functional.interpolate(reference["samples"], size=(latent["samples"].shape[2], latent["samples"].shape[3]), mode="bilinear")
|
|
|
| batch = latent["samples"].shape[0] + reference["samples"].shape[0]
|
| def reference_apply(q, k, v, extra_options):
|
| k = k.clone().repeat(1, 2, 1)
|
| offset = 0
|
| if q.shape[0] > batch:
|
| offset = batch
|
|
|
| for o in range(0, q.shape[0], batch):
|
| for x in range(1, batch):
|
| k[x + o, q.shape[1]:] = q[o,:]
|
|
|
| return q, k, k
|
|
|
| model_reference.set_model_attn1_patch(reference_apply)
|
| out_latent = torch.cat((reference["samples"], latent["samples"]))
|
| if "noise_mask" in latent:
|
| mask = latent["noise_mask"]
|
| else:
|
| mask = torch.ones((64,64), dtype=torch.float32, device="cpu")
|
|
|
| if len(mask.shape) < 3:
|
| mask = mask.unsqueeze(0)
|
| if mask.shape[0] < latent["samples"].shape[0]:
|
| mask = mask.repeat(latent["samples"].shape[0], 1, 1)
|
|
|
| out_mask = torch.zeros((1,mask.shape[1],mask.shape[2]), dtype=torch.float32, device="cpu")
|
| return (model_reference, {"samples": out_latent, "noise_mask": torch.cat((out_mask, mask))})
|
|
|
| NODE_CLASS_MAPPINGS = {
|
| "SwarmReferenceOnly": SwarmReferenceOnly,
|
| }
|
|
|
