|
|
import os |
|
|
|
|
|
import torch |
|
|
|
|
|
from diffusers import DiffusionPipeline, FluxPipeline, StableDiffusion3Pipeline |
|
|
from huggingface_hub import hf_hub_download |
|
|
from safetensors.torch import load_file |
|
|
|
|
|
from .model import DrUM as backbone |
|
|
from .sampling import coreset_sampling |
|
|
|
|
|
def stable_diffusion(large): |
|
|
""" |
|
|
openai/clip-vit-large-patch14, CLIPTextModel, skip -1 |
|
|
""" |
|
|
def inference(prompt, ref_prompt = None, weight = None, alpha = 0.3, skip = -1, batch_size = 64, **kwargs): |
|
|
return large(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, **kwargs), None |
|
|
return inference |
|
|
|
|
|
def stable_diffusion_v2(huge): |
|
|
""" |
|
|
openai/clip-vit-huge-patch14, CLIPTextModel, skip -1 |
|
|
""" |
|
|
def inference(prompt, ref_prompt = None, weight = None, alpha = 0.3, skip = -1, batch_size = 64, **kwargs): |
|
|
return huge(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, **kwargs), None |
|
|
return inference |
|
|
|
|
|
def stable_diffusion_xl(large, bigG): |
|
|
""" |
|
|
openai/clip-vit-large-patch14, CLIPTextModel, skip -2, unnorm |
|
|
laion/CLIP-ViT-bigG-14-laion2B-39B-b160k, CLIPTextModelWithProjection, skip -2, unnorm, pooling + proj |
|
|
""" |
|
|
def inference(prompt, ref_prompt = None, weight = None, alpha = 0.3, skip = -2, batch_size = 64, **kwargs): |
|
|
hidden_state = large(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, normalize = False, **kwargs) |
|
|
if skip == -1: |
|
|
hidden_state2, pool_hidden_state = bigG(prompt, ref_prompt, pooling = True, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, normalize = False, normalize_pool = True, **kwargs) |
|
|
else: |
|
|
hidden_state2 = bigG(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, normalize = False, **kwargs) |
|
|
pool_hidden_state = bigG(prompt, ref_prompt, pooling = True, weight = weight, alpha = alpha, skip = -1, batch_size = batch_size, normalize = False, normalize_pool = True, **kwargs)[1] |
|
|
hidden_state = torch.cat([hidden_state, hidden_state2], dim = -1) |
|
|
pool_hidden_state = bigG.projection_text_hidden_state(pool_hidden_state) |
|
|
return hidden_state.type(pool_hidden_state.dtype), pool_hidden_state |
|
|
return inference |
|
|
|
|
|
def stable_diffusion_v3(large, bigG, t5): |
|
|
""" |
|
|
openai/clip-vit-large-patch14, CLIPTextModelWithProjection, skip -2, unnorm, pooling + proj |
|
|
laion/CLIP-ViT-bigG-14-laion2B-39B-b160k, CLIPTextModelWithProjection, skip -2, unnorm, pooling + proj |
|
|
t5-v1_1-xxl, T5EncoderModel |
|
|
""" |
|
|
def inference(prompt, ref_prompt = None, weight = None, alpha = 0.3, skip = -2, batch_size = 64, **kwargs): |
|
|
if skip == -1: |
|
|
hidden_state, pool_hidden_state = large(prompt, ref_prompt, pooling = True, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, normalize = False, normalize_pool = True, **kwargs) |
|
|
hidden_state2, pool_hidden_state2 = bigG(prompt, ref_prompt, pooling = True, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, normalize = False, normalize_pool = True, **kwargs) |
|
|
else: |
|
|
hidden_state = large(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, normalize = False, **kwargs) |
|
|
hidden_state2 = bigG(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, skip = skip, batch_size = batch_size, normalize = False, **kwargs) |
|
|
pool_hidden_state = large(prompt, ref_prompt, pooling = True, weight = weight, alpha = alpha, skip = -1, batch_size = batch_size, normalize = False, normalize_pool = True, **kwargs)[1] |
|
|
pool_hidden_state2 = bigG(prompt, ref_prompt, pooling = True, weight = weight, alpha = alpha, skip = -1, batch_size = batch_size, normalize = False, normalize_pool = True, **kwargs)[1] |
|
|
hidden_state3 = t5(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, batch_size = batch_size, normalize = False, **kwargs) |
|
|
hidden_state = torch.cat([hidden_state, hidden_state2], dim = -1) |
|
|
pool_hidden_state = large.projection_text_hidden_state(pool_hidden_state) |
|
|
pool_hidden_state2 = bigG.projection_text_hidden_state(pool_hidden_state2) |
|
|
hidden_state = torch.nn.functional.pad(hidden_state, (0, hidden_state3.shape[-1] - hidden_state.shape[-1])) |
|
|
hidden_state = torch.cat([hidden_state, hidden_state3], dim = -2) |
|
|
pool_hidden_state = torch.cat([pool_hidden_state, pool_hidden_state2], dim = -1) |
|
|
return hidden_state.type(pool_hidden_state.dtype), pool_hidden_state |
|
|
return inference |
|
|
|
|
|
def flux(large, t5): |
|
|
""" |
|
|
openai/clip-vit-large-patch14, CLIPTextModel, pooling |
|
|
t5-v1_1-xxl, T5EncoderModel |
|
|
""" |
|
|
def inference(prompt, ref_prompt = None, weight = None, alpha = 0.3, skip = None, batch_size = 64, **kwargs): |
|
|
hidden_state = t5(prompt, ref_prompt, pooling = False, weight = weight, alpha = alpha, batch_size = batch_size, normalize = False, **kwargs) |
|
|
pool_hidden_state = large(prompt, ref_prompt, pooling = True, weight = weight, alpha = alpha, skip = -1, batch_size = batch_size, normalize = False, normalize_pool = True, **kwargs)[1] |
|
|
return hidden_state.type(pool_hidden_state.dtype), pool_hidden_state |
|
|
return inference |
|
|
|
|
|
def peca(pipeline, save_path = "./weight", n_layer = 10): |
|
|
if os.path.exists(os.path.join(save_path, "L.pth")) or os.path.exists(os.path.join(save_path, "H.pth")): |
|
|
load_func = torch.load |
|
|
postfix = "pth" |
|
|
else: |
|
|
from safetensors.torch import load_file as load_func |
|
|
postfix = "safetensors" |
|
|
|
|
|
if "flux" in pipeline.config._name_or_path.split("/")[-1].lower(): |
|
|
model = pipeline.text_encoder |
|
|
processor = pipeline.tokenizer |
|
|
model2 = pipeline.text_encoder_2 |
|
|
processor2 = pipeline.tokenizer_2 |
|
|
|
|
|
large = backbone(model, processor, n_layer = n_layer, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
large.adapter.load_state_dict(load_func(os.path.join(save_path, "L.{0}".format(postfix)))) |
|
|
t5 = backbone(model2, processor2, n_layer = n_layer, encode_ratio = 4, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
t5.adapter.load_state_dict(load_func(os.path.join(save_path, "T5.{0}".format(postfix)))) |
|
|
empty, pool = large.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
large.adapter.set_base_query(torch.cat([empty, pool.unsqueeze(1)], dim = 1)) |
|
|
empty, pool = t5.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
t5.adapter.set_base_query(empty) |
|
|
|
|
|
feature_encoder = large |
|
|
encoder = flux(large, t5) |
|
|
size = 1024 |
|
|
num_inference_steps = 28 |
|
|
skip = -2 |
|
|
elif "stable-diffusion-3.5" in pipeline.config._name_or_path.split("/")[-1].lower(): |
|
|
model = pipeline.text_encoder |
|
|
processor = pipeline.tokenizer |
|
|
model2 = pipeline.text_encoder_2 |
|
|
processor2 = pipeline.tokenizer_2 |
|
|
model3 = pipeline.text_encoder_3 |
|
|
processor3 = pipeline.tokenizer_3 |
|
|
|
|
|
large = backbone(model, processor, n_layer = n_layer, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
large.adapter.load_state_dict(load_func(os.path.join(save_path, "L.{0}".format(postfix)))) |
|
|
bigG = backbone(model2, processor2, n_layer = n_layer, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
bigG.adapter.load_state_dict(load_func(os.path.join(save_path, "bigG.{0}".format(postfix)))) |
|
|
t5 = backbone(model3, processor3, n_layer = n_layer, encode_ratio = 4, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
t5.adapter.load_state_dict(load_func(os.path.join(save_path, "T5.{0}".format(postfix)))) |
|
|
empty, pool = large.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
large.adapter.set_base_query(torch.cat([empty, pool.unsqueeze(1)], dim = 1)) |
|
|
empty, pool = bigG.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
bigG.adapter.set_base_query(torch.cat([empty, pool.unsqueeze(1)], dim = 1)) |
|
|
empty, pool = t5.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
t5.adapter.set_base_query(empty) |
|
|
|
|
|
feature_encoder = large |
|
|
encoder = stable_diffusion_v3(large, bigG, t5) |
|
|
size = 1024 |
|
|
num_inference_steps = 28 |
|
|
skip = -2 |
|
|
elif "xl-base" in pipeline.config._name_or_path.split("/")[-1].lower(): |
|
|
model = pipeline.text_encoder |
|
|
processor = pipeline.tokenizer |
|
|
model2 = pipeline.text_encoder_2 |
|
|
processor2 = pipeline.tokenizer_2 |
|
|
|
|
|
large = backbone(model, processor, n_layer = n_layer, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
large.adapter.load_state_dict(load_func(os.path.join(save_path, "L.{0}".format(postfix)))) |
|
|
bigG = backbone(model2, processor2, n_layer = n_layer, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
bigG.adapter.load_state_dict(load_func(os.path.join(save_path, "bigG.{0}".format(postfix)))) |
|
|
empty, pool = large.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
large.adapter.set_base_query(torch.cat([empty, pool.unsqueeze(1)], dim = 1)) |
|
|
empty, pool = bigG.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
bigG.adapter.set_base_query(torch.cat([empty, pool.unsqueeze(1)], dim = 1)) |
|
|
|
|
|
feature_encoder = large |
|
|
encoder = stable_diffusion_xl(large, bigG) |
|
|
size = 1024 |
|
|
num_inference_steps = 50 |
|
|
skip = -2 |
|
|
elif "stable-diffusion-2" in pipeline.config._name_or_path.split("/")[-1].lower(): |
|
|
model = pipeline.text_encoder |
|
|
processor = pipeline.tokenizer |
|
|
|
|
|
huge = backbone(model, processor, n_layer = n_layer, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
huge.adapter.load_state_dict(load_func(os.path.join(save_path, "H.{0}".format(postfix)))) |
|
|
empty, pool = huge.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
huge.adapter.set_base_query(torch.cat([empty, pool.unsqueeze(1)], dim = 1)) |
|
|
|
|
|
feature_encoder = huge |
|
|
encoder = stable_diffusion_v2(huge) |
|
|
size = 768 |
|
|
num_inference_steps = 50 |
|
|
skip = -1 |
|
|
else: |
|
|
model = pipeline.text_encoder |
|
|
processor = pipeline.tokenizer |
|
|
|
|
|
large = backbone(model, processor, n_layer = n_layer, pos = False, cls_pos = False, dropout = 0.0).to(pipeline.device).eval() |
|
|
large.adapter.load_state_dict(load_func(os.path.join(save_path, "L.{0}".format(postfix)))) |
|
|
empty, pool = large.encode_prompt("", pooling = True, normalize = False, normalize_pool = False) |
|
|
large.adapter.set_base_query(torch.cat([empty, pool.unsqueeze(1)], dim = 1)) |
|
|
|
|
|
feature_encoder = large |
|
|
encoder = stable_diffusion(large) |
|
|
size = 512 |
|
|
num_inference_steps = 50 |
|
|
skip = -1 |
|
|
return encoder, feature_encoder.get_text_feature, size, num_inference_steps, skip |
|
|
|
|
|
class DrUM(DiffusionPipeline): |
|
|
def __init__(self, pipeline, repo_id = "Burf/DrUM", weight = None, torch_dtype = torch.bfloat16, device = "cuda"): |
|
|
""" |
|
|
DrUM for various T2I diffusion models |
|
|
""" |
|
|
self.pipeline = pipeline if not isinstance(pipeline, str) else self.load_pipeline(pipeline, torch_dtype = torch_dtype, device = device) |
|
|
self.repo_id = repo_id |
|
|
|
|
|
self.adapter, self.feature_encoder, self.size, self.num_inference_steps, self.skip = self.load_peca(self.pipeline, repo_id, weight) |
|
|
|
|
|
@classmethod |
|
|
def from_pretrained(cls, pretrained_model_name_or_path, repo_id = "Burf/DrUM", torch_dtype = torch.bfloat16, device = "cuda", weight=None): |
|
|
""" |
|
|
Load DrUM adapter with appropriate pipeline |
|
|
""" |
|
|
pipeline = cls.load_pipeline(pretrained_model_name_or_path, torch_dtype, device) |
|
|
return cls(pipeline = pipeline, repo_id = repo_id, weight = weight, torch_dtype = torch_dtype, device = device) |
|
|
|
|
|
@staticmethod |
|
|
def load_pipeline(model_id, torch_dtype = torch.bfloat16, device = "cuda"): |
|
|
name = model_id.split("/")[-1].lower() |
|
|
if "flux" in name: |
|
|
pipeline = FluxPipeline.from_pretrained(model_id, torch_dtype = torch_dtype) |
|
|
elif "stable-diffusion-3.5" in name: |
|
|
pipeline = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype = torch_dtype) |
|
|
else: |
|
|
pipeline = DiffusionPipeline.from_pretrained(model_id, torch_dtype = torch_dtype) |
|
|
|
|
|
pipeline = pipeline.to(device if torch.cuda.is_available() else "cpu") |
|
|
|
|
|
return pipeline |
|
|
|
|
|
def load_weight(self, pipeline, repo_id = "Burf/DrUM", weight = None): |
|
|
name = pipeline.config._name_or_path.split("/")[-1].lower() |
|
|
|
|
|
weights = [] |
|
|
if "flux" in name: |
|
|
weights = ["L.safetensors", "T5.safetensors"] |
|
|
elif "stable-diffusion-3.5" in name: |
|
|
weights = ["L.safetensors", "bigG.safetensors", "T5.safetensors"] |
|
|
elif "xl-base" in name: |
|
|
weights = ["L.safetensors", "bigG.safetensors"] |
|
|
elif "stable-diffusion-2" in name: |
|
|
weights = ["H.safetensors"] |
|
|
else: |
|
|
weights = ["L.safetensors"] |
|
|
|
|
|
for weight_file in weights: |
|
|
if isinstance(weight, str) and os.path.exists(os.path.join(weight, weight_file)): |
|
|
weight_path = weight |
|
|
break |
|
|
else: |
|
|
safetensor_path = hf_hub_download(repo_id = repo_id, filename = "weight/" + weight_file) |
|
|
weight_path = os.path.dirname(safetensor_path) |
|
|
return weight_path |
|
|
|
|
|
def load_peca(self, pipeline, repo_id = "Burf/DrUM", weight = None): |
|
|
adapter, feature_encoder, size, num_inference_steps, skip = peca(pipeline, save_path = self.load_weight(pipeline, repo_id, weight)) |
|
|
return adapter, feature_encoder, size, num_inference_steps, skip |
|
|
|
|
|
def __call__(self, prompt, ref = None, weight = None, alpha = 0.3, skip = None, sampling = False, seed = 42, |
|
|
size = None, num_inference_steps = None, num_images_per_prompt = 1): |
|
|
""" |
|
|
Generate images using DrUM adapter |
|
|
|
|
|
Args: |
|
|
prompt: Text prompt for generation |
|
|
ref: Reference prompts (list of strings) |
|
|
weight: Weights for reference prompts (list of floats) |
|
|
alpha: Personalization strength (0-1) |
|
|
skip: Text condition axis |
|
|
sampling: Whether to use coreset sampling for reference selection (default: False) |
|
|
seed: Random seed |
|
|
size: Image size |
|
|
num_inference_steps: Inference steps |
|
|
num_images_per_prompt: Number of images to generate |
|
|
|
|
|
Returns: |
|
|
Personalized images (list of PIL Images) |
|
|
""" |
|
|
size = self.size if size is None else size |
|
|
num_inference_steps = self.num_inference_steps if num_inference_steps is None else num_inference_steps |
|
|
skip = self.skip if skip is None else skip |
|
|
|
|
|
if sampling and isinstance(ref, (tuple, list)) and 1 < len(ref): |
|
|
import numpy as np |
|
|
|
|
|
with torch.no_grad(): |
|
|
feature = self.feature_encoder(ref).cpu().float().numpy() |
|
|
|
|
|
indices = coreset_sampling(feature, weight = weight, seed = seed) |
|
|
ref = np.array(ref)[indices].tolist() |
|
|
|
|
|
if isinstance(weight, (tuple, list)) and len(weight) == len(ref): |
|
|
weight = np.array(weight)[indices].tolist() |
|
|
|
|
|
generator = torch.Generator(self.pipeline.device).manual_seed(seed) |
|
|
with torch.no_grad(): |
|
|
cond, pool_cond = self.adapter(prompt, ref, weight = weight, alpha = alpha, skip = skip) |
|
|
|
|
|
pipe_kwargs = { |
|
|
"num_images_per_prompt": num_images_per_prompt, |
|
|
"num_inference_steps": num_inference_steps, |
|
|
"generator": generator, |
|
|
"height": size, |
|
|
"width": size |
|
|
} |
|
|
|
|
|
pipe_kwargs["prompt_embeds"] = cond.type(self.pipeline.dtype) |
|
|
if pool_cond is not None: |
|
|
pipe_kwargs["pooled_prompt_embeds"] = pool_cond.type(self.pipeline.dtype) |
|
|
|
|
|
name = self.pipeline.config._name_or_path.split("/")[-1].lower() |
|
|
if "flux" in name or "stable-diffusion-3" in name: |
|
|
pipe_kwargs["max_sequence_length"] = 256 |
|
|
|
|
|
images = self.pipeline(**pipe_kwargs).images |
|
|
return images |
