adding utils for sliders

utils/clip_util.py

@@ -0,0 +1,142 @@
+from typing import List, Optional
+import math, random, os
+import pandas as pd
+import numpy as np
+import torch
+from tqdm.auto import tqdm
+from sklearn.decomposition import PCA
+
+
+def extract_clip_features(clip, image, encoder):
+    """
+    Extracts feature embeddings from an image using either CLIP or DINOv2 models.
+    
+    Args:
+        clip (torch.nn.Module): The feature extraction model (either CLIP or DINOv2)
+        image (torch.Tensor): Input image tensor normalized according to model requirements
+        encoder (str): Type of encoder to use ('dinov2-small' or 'clip')
+    
+    Returns:
+        torch.Tensor: Feature embeddings extracted from the image
+        
+    Note:
+        - For DINOv2 models, uses the pooled output features
+        - For CLIP models, uses the image features from the vision encoder
+        - The input image should already be properly resized and normalized
+    """
+    # Handle DINOv2 models
+    if 'dino' in encoder:
+        denoised = clip(image)
+        denoised = denoised.pooler_output
+    # Handle CLIP models
+    else:
+        denoised = clip.get_image_features(image)
+    
+    return denoised
+
+@torch.no_grad()
+def compute_clip_pca(
+    diverse_prompts: List[str],
+    pipe,
+    clip_model,
+    clip_processor,
+    device,
+    guidance_scale,
+    params,
+    total_samples = 5000,
+    num_pca_components = 100,
+    batch_size = 10
+    
+) -> torch.Tensor:
+    """
+    Extract CLIP features from generated images based on prompts.
+    
+    Args:
+        diverse_prompts: List of prompts to generate images from
+        model_components: Various model components needed for generation
+        args: Training arguments
+        
+    Returns:
+        Tensor of CLIP principle components
+    """
+    
+    
+    # Calculate how many total batches we need
+    num_batches = math.ceil(total_samples / batch_size)
+    # Randomly sample prompts (with replacement if needed)
+    sampled_prompts_clip = random.choices(diverse_prompts, k=num_batches)
+    
+    clip_features_path = f"{params['savepath_training_images']}/clip_principle_directions.pt"
+    
+    if os.path.exists(clip_features_path):
+        df = pd.read_csv(f"{params['savepath_training_images']}/training_data.csv")
+        prompts_training = list(df.prompt)
+        image_paths = list(df.image_path)
+        return torch.load(clip_features_path).to(device), prompts_training, image_paths
+    
+    os.makedirs(params['savepath_training_images'], exist_ok=True)
+    
+    # Generate images and extract features
+    img_idx = 0
+    clip_features = []
+    image_paths = []
+    prompts_training = []
+    print('Calculating Semantic PCA')
+    
+    for prompt in tqdm(sampled_prompts_clip):
+        if 'max_sequence_length' in params:
+            images = pipe(prompt, 
+                     num_images_per_prompt = batch_size,
+                     num_inference_steps = params['max_denoising_steps'],
+                     guidance_scale=guidance_scale,
+                     max_sequence_length = params['max_sequence_length'],
+                     height = params['height'],
+                     width = params['width'],
+                     ).images
+        else:  
+            images = pipe(prompt, 
+                         num_images_per_prompt = batch_size,
+                         num_inference_steps = params['max_denoising_steps'],
+                         guidance_scale=guidance_scale,
+                         height = params['height'],
+                         width = params['width'],
+                         ).images
+
+        
+        # Process images
+        clip_inputs = clip_processor(images=images, return_tensors="pt", padding=True)
+        pixel_values = clip_inputs['pixel_values'].to(device)
+        
+        # Get image embeddings
+        with torch.no_grad():
+            image_features = clip_model.get_image_features(pixel_values)
+            
+        # Normalize embeddings
+        clip_feats = image_features / image_features.norm(dim=1, keepdim=True)
+        clip_features.append(clip_feats)
+
+        for im in images:
+            image_path = f"{params['savepath_training_images']}/{img_idx}.png"
+            im.save(image_path)
+            image_paths.append(image_path)
+            prompts_training.append(prompt)
+            img_idx += 1
+
+    
+    clip_features = torch.cat(clip_features)
+
+    
+    # Calculate principle components
+    pca = PCA(n_components=num_pca_components)
+    clip_embeds_np = clip_features.float().cpu().numpy()
+    pca.fit(clip_embeds_np)
+    clip_principles = torch.from_numpy(pca.components_).to(device, dtype=pipe.vae.dtype)
+    
+    # Save results
+    torch.save(clip_principles, clip_features_path)
+    pd.DataFrame({
+        'prompt': prompts_training,
+        'image_path': image_paths
+    }).to_csv(f"{params['savepath_training_images']}/training_data.csv", index=False)
+    
+    return clip_principles, prompts_training, image_paths

utils/flux_utils.py

@@ -0,0 +1,404 @@
+import os , torch
+import argparse
+import copy
+import gc
+import itertools
+import logging
+import math
+
+import random
+import shutil
+import warnings
+from contextlib import nullcontext
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.utils.checkpoint
+import transformers
+from accelerate import Accelerator
+from accelerate.logging import get_logger
+from accelerate.utils import DistributedDataParallelKwargs, ProjectConfiguration, set_seed
+from huggingface_hub import create_repo, upload_folder
+from huggingface_hub.utils import insecure_hashlib
+from PIL import Image
+from PIL.ImageOps import exif_transpose
+from torch.utils.data import Dataset
+from torchvision import transforms
+from torchvision.transforms.functional import crop
+from tqdm.auto import tqdm
+from transformers import CLIPTokenizer, PretrainedConfig, T5TokenizerFast
+
+import diffusers
+from diffusers import (
+    AutoencoderKL,
+    FlowMatchEulerDiscreteScheduler,
+    FluxTransformer2DModel,
+)
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import (
+    _set_state_dict_into_text_encoder,
+    cast_training_params,
+    compute_density_for_timestep_sampling,
+    compute_loss_weighting_for_sd3,
+)
+from diffusers.utils import (
+    check_min_version,
+    convert_unet_state_dict_to_peft,
+    is_wandb_available,
+)
+from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
+from diffusers.utils.torch_utils import is_compiled_module
+
+
+from collections import defaultdict
+
+
+from typing import List, Optional
+import argparse
+import ast
+from pathlib import Path
+from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler
+from huggingface_hub import hf_hub_download
+import gc
+import torch.nn.functional as F
+import os
+import torch
+from tqdm.auto import tqdm
+import time, datetime
+import numpy as np
+from torch.optim import AdamW
+from contextlib import ExitStack
+from safetensors.torch import load_file
+import torch.nn as nn
+import random
+from transformers import CLIPModel
+
+from transformers import logging
+logging.set_verbosity_warning()
+
+from diffusers import logging
+logging.set_verbosity_error()
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+flush()
+def unwrap_model(model):
+    options = (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)
+    #if is_deepspeed_available():
+    #    options += (DeepSpeedEngine,)
+    while isinstance(model, options):
+        model = model.module
+    return model
+
+
+# Function to log gradients
+def log_gradients(named_parameters):
+    grad_dict = defaultdict(lambda: defaultdict(float))
+    for name, param in named_parameters:
+        if param.requires_grad and param.grad is not None:
+            grad_dict[name]['mean'] = param.grad.abs().mean().item()
+            grad_dict[name]['std'] = param.grad.std().item()
+            grad_dict[name]['max'] = param.grad.abs().max().item()
+            grad_dict[name]['min'] = param.grad.abs().min().item()
+    return grad_dict
+    
+def import_model_class_from_model_name_or_path(
+    pretrained_model_name_or_path: str, subfolder: str = "text_encoder",
+):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path, subfolder=subfolder
+        , device_map='cuda:0'
+    )
+    model_class = text_encoder_config.architectures[0]
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+
+        return CLIPTextModel
+    elif model_class == "T5EncoderModel":
+        from transformers import T5EncoderModel
+
+        return T5EncoderModel
+    else:
+        raise ValueError(f"{model_class} is not supported.")
+def load_text_encoders(pretrained_model_name_or_path, class_one, class_two, weight_dtype):
+    text_encoder_one = class_one.from_pretrained(
+        pretrained_model_name_or_path, 
+        subfolder="text_encoder", 
+        torch_dtype=weight_dtype,
+        device_map='cuda:0'
+    )
+    text_encoder_two = class_two.from_pretrained(
+        pretrained_model_name_or_path, 
+        subfolder="text_encoder_2", 
+        torch_dtype=weight_dtype,
+        device_map='cuda:0'
+    )
+    return text_encoder_one, text_encoder_two
+import matplotlib.pyplot as plt
+def plot_labeled_images(images, labels):
+    # Determine the number of images
+    n = len(images)
+    
+    # Create a new figure with a single row
+    fig, axes = plt.subplots(1, n, figsize=(5*n, 5))
+    
+    # If there's only one image, axes will be a single object, not an array
+    if n == 1:
+        axes = [axes]
+    
+    # Plot each image
+    for i, (img, label) in enumerate(zip(images, labels)):
+        # Convert PIL image to numpy array
+        img_array = np.array(img)
+        
+        # Display the image
+        axes[i].imshow(img_array)
+        axes[i].axis('off')  # Turn off axis
+        
+        # Set the title (label) for the image
+        axes[i].set_title(label)
+    
+    # Adjust the layout and display the plot
+    plt.tight_layout()
+    plt.show()
+    
+
+def tokenize_prompt(tokenizer, prompt, max_sequence_length):
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=max_sequence_length,
+        truncation=True,
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+    return text_input_ids
+
+
+def _encode_prompt_with_t5(
+    text_encoder,
+    tokenizer,
+    max_sequence_length=512,
+    prompt=None,
+    num_images_per_prompt=1,
+    device=None,
+    text_input_ids=None,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    if tokenizer is not None:
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=max_sequence_length,
+            truncation=True,
+            return_length=False,
+            return_overflowing_tokens=False,
+            return_tensors="pt",
+        )
+        text_input_ids = text_inputs.input_ids
+    else:
+        if text_input_ids is None:
+            raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
+
+    prompt_embeds = text_encoder(text_input_ids.to(device))[0]
+
+    dtype = text_encoder.dtype
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+    _, seq_len, _ = prompt_embeds.shape
+
+    # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+    return prompt_embeds
+
+
+def _encode_prompt_with_clip(
+    text_encoder,
+    tokenizer,
+    prompt: str,
+    device=None,
+    text_input_ids=None,
+    num_images_per_prompt: int = 1,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    if tokenizer is not None:
+        text_inputs = tokenizer(
+            prompt,
+            padding="max_length",
+            max_length=77,
+            truncation=True,
+            return_overflowing_tokens=False,
+            return_length=False,
+            return_tensors="pt",
+        )
+
+        text_input_ids = text_inputs.input_ids
+    else:
+        if text_input_ids is None:
+            raise ValueError("text_input_ids must be provided when the tokenizer is not specified")
+
+    prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+    # Use pooled output of CLIPTextModel
+    prompt_embeds = prompt_embeds.pooler_output
+    prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device)
+
+    # duplicate text embeddings for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+    return prompt_embeds
+
+def encode_prompt(
+    text_encoders,
+    tokenizers,
+    prompt: str,
+    max_sequence_length,
+    device=None,
+    num_images_per_prompt: int = 1,
+    text_input_ids_list=None,
+):
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+    dtype = text_encoders[0].dtype
+
+    pooled_prompt_embeds = _encode_prompt_with_clip(
+        text_encoder=text_encoders[0],
+        tokenizer=tokenizers[0],
+        prompt=prompt,
+        device=device if device is not None else text_encoders[0].device,
+        num_images_per_prompt=num_images_per_prompt,
+        text_input_ids=text_input_ids_list[0] if text_input_ids_list else None,
+    )
+
+    prompt_embeds = _encode_prompt_with_t5(
+        text_encoder=text_encoders[1],
+        tokenizer=tokenizers[1],
+        max_sequence_length=max_sequence_length,
+        prompt=prompt,
+        num_images_per_prompt=num_images_per_prompt,
+        device=device if device is not None else text_encoders[1].device,
+        text_input_ids=text_input_ids_list[1] if text_input_ids_list else None,
+    )
+
+    text_ids = torch.zeros(batch_size, prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+    text_ids = text_ids.repeat(num_images_per_prompt, 1, 1)
+
+    return prompt_embeds, pooled_prompt_embeds, text_ids
+    
+def compute_text_embeddings(prompt, text_encoders, tokenizers,max_sequence_length=256):
+    device = text_encoders[0].device
+    with torch.no_grad():
+        prompt_embeds, pooled_prompt_embeds, text_ids = encode_prompt(
+            text_encoders, tokenizers, prompt, max_sequence_length=max_sequence_length
+        )
+        prompt_embeds = prompt_embeds.to(device)
+        pooled_prompt_embeds = pooled_prompt_embeds.to(device)
+        text_ids = text_ids.to(device)
+    return prompt_embeds, pooled_prompt_embeds, text_ids
+
+
+def get_sigmas(timesteps, n_dim=4, device='cuda:0', dtype=torch.bfloat16):
+    sigmas = noise_scheduler_copy.sigmas.to(device=device, dtype=dtype)
+    schedule_timesteps = noise_scheduler_copy.timesteps.to(device)
+    timesteps = timesteps.to(device)
+    step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]
+
+    sigma = sigmas[step_indices].flatten()
+    while len(sigma.shape) < n_dim:
+        sigma = sigma.unsqueeze(-1)
+    return sigma
+
+
+def plot_history(history):
+    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 5))
+    ax1.plot(history['concept'])
+    ax1.set_title('Concept Loss')
+    ax2.plot(movingaverage(history['concept'], 10))
+    ax2.set_title('Moving Average Concept Loss')
+    plt.tight_layout()
+    plt.show()
+
+def movingaverage(interval, window_size):
+    window = np.ones(int(window_size))/float(window_size)
+    return np.convolve(interval, window, 'same')
+
+
+
+@torch.no_grad()
+def get_noisy_image_flux(
+    image,
+    vae,
+    transformer,
+    scheduler,
+    timesteps_to=1000,
+    generator=None,
+    **kwargs,
+):
+    """
+    Gets noisy latents for a given image using Flux pipeline approach.
+    
+    Args:
+        image: PIL image or tensor
+        vae: Flux VAE model
+        transformer: Flux transformer model
+        scheduler: Flux noise scheduler
+        timesteps_to: Target timestep
+        generator: Random generator for reproducibility
+        
+    Returns:
+        tuple: (noisy_latents, noise)
+    """
+    device = vae.device
+    vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
+
+    # Preprocess image
+    if not isinstance(image, torch.Tensor):
+        image = image_processor.preprocess(image)
+    image = image.to(device=device, dtype=torch.float32)
+
+    # Encode through VAE
+    init_latents = vae.encode(image).latents
+    init_latents = vae.config.scaling_factor * init_latents
+    
+    # Get shape for noise
+    shape = init_latents.shape
+    
+    # Generate noise
+    noise = randn_tensor(shape, generator=generator, device=device)
+
+    # Pack latents using Flux's method
+    init_latents = _pack_latents(
+        init_latents, 
+        shape[0],  # batch size
+        transformer.config.in_channels // 4,
+        height=shape[2],
+        width=shape[3]
+    )
+    noise = _pack_latents(
+        noise,
+        shape[0],
+        transformer.config.in_channels // 4,
+        height=shape[2],
+        width=shape[3]
+    )
+
+    # Get timestep
+    timestep = scheduler.timesteps[timesteps_to:timesteps_to+1]
+    
+    # Add noise to latents
+    noisy_latents = scheduler.add_noise(init_latents, noise, timestep)
+    
+    return noisy_latents, noise

utils/lora.py

@@ -0,0 +1,320 @@
+# ref:
+# - https://github.com/cloneofsimo/lora/blob/master/lora_diffusion/lora.py
+# - https://github.com/kohya-ss/sd-scripts/blob/main/networks/lora.py
+
+import os
+import math
+from typing import Optional, List, Type, Set, Literal
+
+import torch
+import torch.nn as nn
+from diffusers import UNet2DConditionModel
+from safetensors.torch import save_file
+from datetime import datetime
+
+UNET_TARGET_REPLACE_MODULE_TRANSFORMER = [
+#     "Transformer2DModel",  # どうやらこっちの方らしい？ # attn1, 2
+    "Attention"
+]
+UNET_TARGET_REPLACE_MODULE_CONV = [
+    "ResnetBlock2D",
+    "Downsample2D",
+    "Upsample2D",
+    "DownBlock2D",
+    "UpBlock2D",
+    
+]  # locon, 3clier
+
+LORA_PREFIX_UNET = "lora_unet"
+
+DEFAULT_TARGET_REPLACE = UNET_TARGET_REPLACE_MODULE_TRANSFORMER
+
+TRAINING_METHODS = Literal[
+    "noxattn",  # train all layers except x-attns and time_embed layers
+    "innoxattn",  # train all layers except self attention layers
+    "selfattn",  # ESD-u, train only self attention layers
+    "xattn",  # ESD-x, train only x attention layers
+    "xattn-up", # all up blocks only
+    "xattn-down",# all down blocks only
+    "xattn-mid",# mid blocks only
+    "full",  #  train all layers
+    "xattn-strict", # q and k values
+    "noxattn-hspace",
+    "noxattn-hspace-last",
+    "flux-attn",
+    # "xlayer",
+    # "outxattn",
+    # "outsattn",
+    # "inxattn",
+    # "inmidsattn",
+    # "selflayer",
+]
+
+def load_ortho_dict(n):
+    path = f'/share/u/rohit/orthogonal_basis/{n:09}.ckpt'
+    if os.path.isfile(path):
+        return torch.load(path)
+    else:
+        x = torch.randn(n,n)
+        eig, _, _ = torch.svd(x)
+        torch.save(eig, path)
+        return eig
+
+def init_ortho_proj(rank, weight):
+    seed = torch.seed()
+    torch.manual_seed(datetime.now().timestamp())
+    q_index = torch.randint(high=weight.size(0),size=(rank,))
+    torch.manual_seed(seed)
+
+    ortho_q_init = load_ortho_dict(weight.size(0)).to(dtype=weight.dtype)[:,q_index]
+    return nn.Parameter(ortho_q_init)
+
+
+class LoRAModule(nn.Module):
+    """
+    replaces forward method of the original Linear, instead of replacing the original Linear module.
+    """
+
+    def __init__(
+        self,
+        lora_name,
+        org_module: nn.Module,
+        multiplier=1.0,
+        lora_dim=4,
+        alpha=1,
+        train_method='xattn',
+        fast_init = False
+    ):
+        """if alpha == 0 or None, alpha is rank (no scaling)."""
+        super().__init__()
+        self.lora_name = lora_name
+        self.lora_dim = lora_dim
+
+        if "Linear" in org_module.__class__.__name__:
+            in_dim = org_module.in_features
+            out_dim = org_module.out_features
+            self.lora_down = nn.Linear(in_dim, lora_dim, bias=False)
+            self.lora_up = nn.Linear(lora_dim, out_dim, bias=False)
+
+        elif "Conv" in org_module.__class__.__name__:  # 一応
+            in_dim = org_module.in_channels
+            out_dim = org_module.out_channels
+
+            self.lora_dim = min(self.lora_dim, in_dim, out_dim)
+            if self.lora_dim != lora_dim:
+                print(f"{lora_name} dim (rank) is changed to: {self.lora_dim}")
+
+            kernel_size = org_module.kernel_size
+            stride = org_module.stride
+            padding = org_module.padding
+            self.lora_down = nn.Conv2d(
+                in_dim, self.lora_dim, kernel_size, stride, padding, bias=False
+            )
+            self.lora_up = nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
+
+        if type(alpha) == torch.Tensor:
+            alpha = alpha.detach().numpy()
+        alpha = lora_dim if alpha is None or alpha == 0 else alpha
+        self.scale = alpha / self.lora_dim
+        self.register_buffer("alpha", torch.tensor(alpha))  # 定数として扱える
+
+        # same as microsoft's
+        nn.init.kaiming_uniform_(self.lora_down.weight, a=1)
+        if train_method == 'full':
+            nn.init.zeros_(self.lora_up.weight)
+        else:
+            if not fast_init:
+                self.lora_up.weight = init_ortho_proj(lora_dim, self.lora_up.weight)
+                self.lora_up.weight.requires_grad_(False)
+            else:
+                nn.init.zeros_(self.lora_up.weight)
+        
+        self.multiplier = multiplier
+        self.org_module = org_module  # remove in applying
+
+    def apply_to(self):
+        self.org_forward = self.org_module.forward
+        self.org_module.forward = self.forward
+        del self.org_module
+
+    def forward(self, x):
+        return (
+            self.org_forward(x)
+            + self.lora_up(self.lora_down(x)) * self.multiplier * self.scale
+        )
+
+
+class LoRANetwork(nn.Module):
+    def __init__(
+        self,
+        unet: UNet2DConditionModel,
+        rank: int = 4,
+        multiplier: float = 1.0,
+        alpha: float = 1.0,
+        train_method: TRAINING_METHODS = "full",
+        layers = ['Linear', 'Conv'],
+        fast_init = False,
+    ) -> None:
+        super().__init__()
+        self.lora_scale = 1
+        self.multiplier = multiplier
+        self.lora_dim = rank
+        self.alpha = alpha
+        self.train_method=train_method
+        # LoRAのみ
+        self.module = LoRAModule
+
+        # unetのloraを作る
+        self.unet_loras = self.create_modules(
+            LORA_PREFIX_UNET,
+            unet,
+            DEFAULT_TARGET_REPLACE,
+            self.lora_dim,
+            self.multiplier,
+            train_method=train_method,
+            layers = layers,
+            fast_init=fast_init,
+        )
+        print(f"create LoRA for U-Net: {len(self.unet_loras)} modules.")
+
+        # assertion 名前の被りがないか確認しているようだ
+        lora_names = set()
+        for lora in self.unet_loras:
+            assert (
+                lora.lora_name not in lora_names
+            ), f"duplicated lora name: {lora.lora_name}. {lora_names}"
+            lora_names.add(lora.lora_name)
+
+        # 適用する
+        for lora in self.unet_loras:
+            lora.apply_to()
+            self.add_module(
+                lora.lora_name,
+                lora,
+            )
+
+        del unet
+
+        torch.cuda.empty_cache()
+
+    def create_modules(
+        self,
+        prefix: str,
+        root_module: nn.Module,
+        target_replace_modules: List[str],
+        rank: int,
+        multiplier: float,
+        train_method: TRAINING_METHODS,
+        layers: List[str],
+        fast_init: bool,
+    ) -> list:
+        filt_layers = []
+        if 'Linear' in layers:
+            filt_layers.extend(["Linear", "LoRACompatibleLinear"])
+        if 'Conv' in layers:
+            filt_layers.extend(["Conv2d", "LoRACompatibleConv"])
+        loras = []
+        names = []
+        for name, module in root_module.named_modules():
+            if train_method == "noxattn" or train_method == "noxattn-hspace" or train_method == "noxattn-hspace-last":  # Cross Attention と Time Embed 以外学習
+                if "attn2" in name or "time_embed" in name:
+                    continue
+            elif train_method == "innoxattn":  # Cross Attention 以外学習
+                if "attn2" in name:
+                    continue
+            elif train_method == "selfattn":  # Self Attention のみ学習
+                if "attn1" not in name:
+                    continue
+            elif train_method in ["xattn", "xattn-strict", "xattn-up", "xattn-down", "xattn-mid"]:  # Cross Attention のみ学習
+                if "attn2" not in name:
+                    continue
+                if train_method == 'xattn-up':
+                    if 'up_block' not in name:
+                        continue
+                if train_method == 'xattn-down':
+                    if 'down_block' not in name:
+                        continue
+                if train_method == 'xattn-mid':
+                    if 'mid_block' not in name:
+                        continue
+            elif train_method == "full":  # 全部学習
+                pass
+            elif train_method == "flux-attn":
+                if "attn" not in name:
+                    continue
+            else:
+                raise NotImplementedError(
+                    f"train_method: {train_method} is not implemented."
+                )
+            if module.__class__.__name__ in target_replace_modules:
+                for child_name, child_module in module.named_modules():
+                    if child_module.__class__.__name__ in filt_layers:
+                        
+                            
+                        if train_method == 'xattn-strict':
+                            if 'out' in child_name:
+                                continue
+                            if 'to_q' in child_name:
+                                continue
+                        if train_method == 'noxattn-hspace':
+                            if 'mid_block' not in name:
+                                continue
+                        if train_method == 'noxattn-hspace-last':
+                            if 'mid_block' not in name or '.1' not in name or 'conv2' not in child_name:
+                                continue
+                        lora_name = prefix + "." + name + "." + child_name
+                        lora_name = lora_name.replace(".", "_")
+#                         print(f"{lora_name}")
+                        lora = self.module(
+                            lora_name, child_module, multiplier, rank, self.alpha, train_method, fast_init
+                        )
+#                         print(name, child_name)
+#                         print(child_module.weight.shape)
+                        if lora_name not in names:
+                            loras.append(lora)
+                            names.append(lora_name)
+        # print(f'@@@@@@@@@@@@@@@@@@@@@@@@@@@@ \n {names}')
+        return loras
+
+    def prepare_optimizer_params(self):
+        all_params = []
+
+        if self.unet_loras:  # 実質これしかない
+            params = []
+            if self.train_method == 'full':
+                [params.extend(lora.parameters()) for lora in self.unet_loras]
+            else:
+                [params.extend(lora.lora_down.parameters()) for lora in self.unet_loras]
+            param_data = {"params": params}
+            all_params.append(param_data)
+
+        return all_params
+
+    def save_weights(self, file, dtype=None, metadata: Optional[dict] = None):
+        state_dict = self.state_dict()
+
+        if dtype is not None:
+            for key in list(state_dict.keys()):
+                v = state_dict[key]
+                v = v.detach().clone().to("cpu").to(dtype)
+                state_dict[key] = v
+
+#         for key in list(state_dict.keys()):
+#             if not key.startswith("lora"):
+#                 # lora以外除外
+#                 del state_dict[key]
+
+        if os.path.splitext(file)[1] == ".safetensors":
+            save_file(state_dict, file, metadata)
+        else:
+            torch.save(state_dict, file)
+    def set_lora_slider(self, scale):
+        self.lora_scale = scale
+
+    def __enter__(self):
+        for lora in self.unet_loras:
+            lora.multiplier = 1.0 * self.lora_scale
+
+    def __exit__(self, exc_type, exc_value, tb):
+        for lora in self.unet_loras:
+            lora.multiplier = 0

utils/model_util.py

@@ -0,0 +1,527 @@
+from typing import Literal, Union, Optional
+
+import torch, gc, os
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextModelWithProjection, T5TokenizerFast
+from transformers import (
+    AutoModel,
+    CLIPModel,
+    CLIPProcessor,
+)
+from huggingface_hub import hf_hub_download
+from diffusers import (
+    UNet2DConditionModel,
+    SchedulerMixin,
+    StableDiffusionPipeline,
+    StableDiffusionXLPipeline,
+    FluxPipeline,
+    AutoencoderKL,
+    FluxTransformer2DModel,
+)
+import copy
+from diffusers.schedulers import (
+    DDIMScheduler,
+    DDPMScheduler,
+    LMSDiscreteScheduler,
+    EulerAncestralDiscreteScheduler,
+    FlowMatchEulerDiscreteScheduler,
+)
+from diffusers import LCMScheduler, AutoencoderTiny
+import sys
+sys.path.append('.')
+from .flux_utils import *
+
+TOKENIZER_V1_MODEL_NAME = "CompVis/stable-diffusion-v1-4"
+TOKENIZER_V2_MODEL_NAME = "stabilityai/stable-diffusion-2-1"
+
+AVAILABLE_SCHEDULERS = Literal["ddim", "ddpm", "lms", "euler_a"]
+
+SDXL_TEXT_ENCODER_TYPE = Union[CLIPTextModel, CLIPTextModelWithProjection]
+
+DIFFUSERS_CACHE_DIR = None  # if you want to change the cache dir, change this
+
+
+def load_diffusers_model(
+    pretrained_model_name_or_path: str,
+    v2: bool = False,
+    clip_skip: Optional[int] = None,
+    weight_dtype: torch.dtype = torch.float32,
+) -> tuple[CLIPTokenizer, CLIPTextModel, UNet2DConditionModel,]:
+    # VAE はいらない
+
+    if v2:
+        tokenizer = CLIPTokenizer.from_pretrained(
+            TOKENIZER_V2_MODEL_NAME,
+            subfolder="tokenizer",
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+        )
+        text_encoder = CLIPTextModel.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="text_encoder",
+            # default is clip skip 2
+            num_hidden_layers=24 - (clip_skip - 1) if clip_skip is not None else 23,
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+        )
+    else:
+        tokenizer = CLIPTokenizer.from_pretrained(
+            TOKENIZER_V1_MODEL_NAME,
+            subfolder="tokenizer",
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+        )
+        text_encoder = CLIPTextModel.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="text_encoder",
+            num_hidden_layers=12 - (clip_skip - 1) if clip_skip is not None else 12,
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+        )
+
+    unet = UNet2DConditionModel.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="unet",
+        torch_dtype=weight_dtype,
+        cache_dir=DIFFUSERS_CACHE_DIR,
+    )
+
+    return tokenizer, text_encoder, unet
+
+
+def load_checkpoint_model(
+    checkpoint_path: str,
+    v2: bool = False,
+    clip_skip: Optional[int] = None,
+    weight_dtype: torch.dtype = torch.float32,
+) -> tuple[CLIPTokenizer, CLIPTextModel, UNet2DConditionModel,]:
+    pipe = StableDiffusionPipeline.from_ckpt(
+        checkpoint_path,
+        upcast_attention=True if v2 else False,
+        torch_dtype=weight_dtype,
+        cache_dir=DIFFUSERS_CACHE_DIR,
+    )
+
+    unet = pipe.unet
+    tokenizer = pipe.tokenizer
+    text_encoder = pipe.text_encoder
+    if clip_skip is not None:
+        if v2:
+            text_encoder.config.num_hidden_layers = 24 - (clip_skip - 1)
+        else:
+            text_encoder.config.num_hidden_layers = 12 - (clip_skip - 1)
+
+    del pipe
+
+    return tokenizer, text_encoder, unet
+
+
+def load_models(
+    pretrained_model_name_or_path: str,
+    scheduler_name: AVAILABLE_SCHEDULERS,
+    v2: bool = False,
+    v_pred: bool = False,
+    weight_dtype: torch.dtype = torch.float32,
+) -> tuple[CLIPTokenizer, CLIPTextModel, UNet2DConditionModel, SchedulerMixin,]:
+    if pretrained_model_name_or_path.endswith(
+        ".ckpt"
+    ) or pretrained_model_name_or_path.endswith(".safetensors"):
+        tokenizer, text_encoder, unet = load_checkpoint_model(
+            pretrained_model_name_or_path, v2=v2, weight_dtype=weight_dtype
+        )
+    else:  # diffusers
+        tokenizer, text_encoder, unet = load_diffusers_model(
+            pretrained_model_name_or_path, v2=v2, weight_dtype=weight_dtype
+        )
+
+    # VAE はいらない
+
+    scheduler = create_noise_scheduler(
+        scheduler_name,
+        prediction_type="v_prediction" if v_pred else "epsilon",
+    )
+
+    return tokenizer, text_encoder, unet, scheduler
+
+
+def load_diffusers_model_xl(
+    pretrained_model_name_or_path: str,
+    weight_dtype: torch.dtype = torch.float32,
+) -> tuple[list[CLIPTokenizer], list[SDXL_TEXT_ENCODER_TYPE], UNet2DConditionModel,]:
+    # returns tokenizer, tokenizer_2, text_encoder, text_encoder_2, unet
+
+    tokenizers = [
+        CLIPTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="tokenizer",
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+        ),
+        CLIPTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="tokenizer_2",
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+            pad_token_id=0,  # same as open clip
+        ),
+    ]
+
+    text_encoders = [
+        CLIPTextModel.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="text_encoder",
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+        ),
+        CLIPTextModelWithProjection.from_pretrained(
+            pretrained_model_name_or_path,
+            subfolder="text_encoder_2",
+            torch_dtype=weight_dtype,
+            cache_dir=DIFFUSERS_CACHE_DIR,
+        ),
+    ]
+
+    unet = UNet2DConditionModel.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="unet",
+        torch_dtype=weight_dtype,
+        cache_dir=DIFFUSERS_CACHE_DIR,
+    )
+
+    return tokenizers, text_encoders, unet
+
+
+def load_checkpoint_model_xl(
+    checkpoint_path: str,
+    weight_dtype: torch.dtype = torch.float32,
+) -> tuple[list[CLIPTokenizer], list[SDXL_TEXT_ENCODER_TYPE], UNet2DConditionModel,]:
+    pipe = StableDiffusionXLPipeline.from_single_file(
+        checkpoint_path,
+        torch_dtype=weight_dtype,
+        cache_dir=DIFFUSERS_CACHE_DIR,
+    )
+
+    unet = pipe.unet
+    tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
+    text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
+    if len(text_encoders) == 2:
+        text_encoders[1].pad_token_id = 0
+
+    del pipe
+
+    return tokenizers, text_encoders, unet
+
+
+def load_models_xl_(
+    pretrained_model_name_or_path: str,
+    scheduler_name: AVAILABLE_SCHEDULERS,
+    weight_dtype: torch.dtype = torch.float32,
+) -> tuple[
+    list[CLIPTokenizer],
+    list[SDXL_TEXT_ENCODER_TYPE],
+    UNet2DConditionModel,
+    SchedulerMixin,
+]:
+    if pretrained_model_name_or_path.endswith(
+        ".ckpt"
+    ) or pretrained_model_name_or_path.endswith(".safetensors"):
+        (
+            tokenizers,
+            text_encoders,
+            unet,
+        ) = load_checkpoint_model_xl(pretrained_model_name_or_path, weight_dtype)
+    else:  # diffusers
+        (
+            tokenizers,
+            text_encoders,
+            unet,
+        ) = load_diffusers_model_xl(pretrained_model_name_or_path, weight_dtype)
+
+    scheduler = create_noise_scheduler(scheduler_name)
+        
+    return tokenizers, text_encoders, unet, scheduler
+
+
+def create_noise_scheduler(
+    scheduler_name: AVAILABLE_SCHEDULERS = "ddpm",
+    prediction_type: Literal["epsilon", "v_prediction"] = "epsilon",
+) -> SchedulerMixin:
+    # 正直、どれがいいのかわからない。元の実装だとDDIMとDDPMとLMSを選べたのだけど、どれがいいのかわからぬ。
+
+    name = scheduler_name.lower().replace(" ", "_")
+    if name == "ddim":
+        # https://huggingface.co/docs/diffusers/v0.17.1/en/api/schedulers/ddim
+        scheduler = DDIMScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000,
+            clip_sample=False,
+            prediction_type=prediction_type,  # これでいいの？
+        )
+    elif name == "ddpm":
+        # https://huggingface.co/docs/diffusers/v0.17.1/en/api/schedulers/ddpm
+        scheduler = DDPMScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000,
+            clip_sample=False,
+            prediction_type=prediction_type,
+        )
+    elif name == "lms":
+        # https://huggingface.co/docs/diffusers/v0.17.1/en/api/schedulers/lms_discrete
+        scheduler = LMSDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000,
+            prediction_type=prediction_type,
+        )
+    elif name == "euler_a":
+        # https://huggingface.co/docs/diffusers/v0.17.1/en/api/schedulers/euler_ancestral
+        scheduler = EulerAncestralDiscreteScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000,
+            # clip_sample=False,
+            prediction_type=prediction_type,
+        )
+    else:
+        raise ValueError(f"Unknown scheduler name: {name}")
+
+    return scheduler
+
+
+def load_models_xl(params):
+    """
+    Load all required models for training
+    
+    Args:
+        params: Dictionary containing model parameters and configurations
+        
+    Returns:
+        dict: Dictionary containing all loaded models and tokenizers
+    """
+    device = params['device']
+    weight_dtype = params['weight_dtype']
+    
+    # Load SDXL components (UNet, text encoders, tokenizers)
+    scheduler_name = 'ddim'
+    tokenizers, text_encoders, unet, noise_scheduler = load_models_xl_(
+        params['pretrained_model_name_or_path'],
+        scheduler_name=scheduler_name,
+    )
+    
+    # Move text encoders to device and set to eval mode
+    for text_encoder in text_encoders:
+        text_encoder.to(device, dtype=weight_dtype)
+        text_encoder.requires_grad_(False)
+        text_encoder.eval()
+    
+    # Set up UNet
+    unet.to(device, dtype=weight_dtype)
+    unet.requires_grad_(False)
+    unet.eval()
+    
+    # Load tiny VAE for efficiency
+    vae = AutoencoderTiny.from_pretrained(
+        "madebyollin/taesdxl", 
+        torch_dtype=weight_dtype
+    )
+    vae = vae.to(device, dtype=weight_dtype)
+    vae.requires_grad_(False)
+    
+    # Load appropriate encoder (CLIP or DinoV2)
+    if params['encoder'] == 'dinov2-small':
+        clip_model = AutoModel.from_pretrained(
+            'facebook/dinov2-small', 
+            torch_dtype=weight_dtype
+        )
+        clip_processor= None
+    else:
+        clip_model = CLIPModel.from_pretrained(
+            "wkcn/TinyCLIP-ViT-40M-32-Text-19M-LAION400M", 
+            torch_dtype=weight_dtype
+        )
+        clip_processor = CLIPProcessor.from_pretrained("wkcn/TinyCLIP-ViT-40M-32-Text-19M-LAION400M")
+    clip_model = clip_model.to(device, dtype=weight_dtype)
+    clip_model.requires_grad_(False)
+
+    
+    
+    # If using DMD checkpoint, load it
+    if params['distilled'] != 'None':
+        if '.safetensors' in params['distilled']:
+            unet.load_state_dict(load_file(params['distilled'], device=device))
+        elif 'dmd2' in params['distilled']:
+            repo_name = "tianweiy/DMD2"
+            ckpt_name = "dmd2_sdxl_4step_unet_fp16.bin"
+            unet.load_state_dict(torch.load(hf_hub_download(repo_name, ckpt_name)))
+        else:
+            unet.load_state_dict(torch.load(params['distilled']))
+
+        
+        # Set up LCM scheduler for DMD
+        noise_scheduler = LCMScheduler(
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            num_train_timesteps=1000,
+            prediction_type="epsilon",
+            original_inference_steps=1000
+        )
+
+    noise_scheduler.set_timesteps(params['max_denoising_steps'])
+    pipe = StableDiffusionXLPipeline(vae = vae,
+            text_encoder = text_encoders[0],
+            text_encoder_2 = text_encoders[1],
+            tokenizer = tokenizers[0],
+            tokenizer_2 = tokenizers[1],
+            unet = unet,
+            scheduler = noise_scheduler)
+    pipe.set_progress_bar_config(disable=True)
+    return {
+        'unet': unet,
+        'vae': vae,
+        'clip_model': clip_model,
+        'clip_processor': clip_processor,
+        'tokenizers': tokenizers,
+        'text_encoders': text_encoders,
+        'noise_scheduler': noise_scheduler
+    }, pipe
+
+
+def load_models_flux(params):
+    # Load the tokenizers
+    tokenizer_one = CLIPTokenizer.from_pretrained(
+        params['pretrained_model_name_or_path'],
+        subfolder="tokenizer",
+        torch_dtype=params['weight_dtype'], device_map=params['device']
+    )
+    tokenizer_two = T5TokenizerFast.from_pretrained(
+        params['pretrained_model_name_or_path'],
+        subfolder="tokenizer_2",
+        torch_dtype=params['weight_dtype'], device_map=params['device']
+    )
+    
+    # Load scheduler and models
+    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+            params['pretrained_model_name_or_path'], 
+            subfolder="scheduler",
+            torch_dtype=params['weight_dtype'], device=params['device']
+    )
+    noise_scheduler_copy = copy.deepcopy(noise_scheduler)
+    
+    
+    
+    # import correct text encoder classes
+    text_encoder_cls_one = import_model_class_from_model_name_or_path(
+        params['pretrained_model_name_or_path'],
+    )
+    text_encoder_cls_two = import_model_class_from_model_name_or_path(
+       params['pretrained_model_name_or_path'], subfolder="text_encoder_2"
+    )
+    # Load the text encoders
+    text_encoder_one, text_encoder_two = load_text_encoders(params['pretrained_model_name_or_path'], text_encoder_cls_one, text_encoder_cls_two, params['weight_dtype'])
+    
+    # Load VAE
+    vae = AutoencoderKL.from_pretrained(
+        params['pretrained_model_name_or_path'],
+        subfolder="vae",
+        torch_dtype=params['weight_dtype'], device_map='auto'
+    )
+    transformer = FluxTransformer2DModel.from_pretrained(
+        params['pretrained_model_name_or_path'], 
+        subfolder="transformer", 
+        torch_dtype=params['weight_dtype']
+    )
+    
+    # We only train the additional adapter LoRA layers
+    transformer.requires_grad_(False)
+    vae.requires_grad_(False)
+    text_encoder_one.requires_grad_(False)
+    text_encoder_two.requires_grad_(False)
+    
+    vae.to(params['device'])
+    transformer.to(params['device'])
+    text_encoder_one.to(params['device'])
+    text_encoder_two.to(params['device'])
+
+    # Load appropriate encoder (CLIP or DinoV2)
+    if params['encoder'] == 'dinov2-small':
+        clip_model = AutoModel.from_pretrained(
+            'facebook/dinov2-small', 
+            torch_dtype=params['weight_dtype']
+        )
+        clip_processor= None
+    else:
+        clip_model = CLIPModel.from_pretrained(
+            "wkcn/TinyCLIP-ViT-40M-32-Text-19M-LAION400M", 
+            torch_dtype=params['weight_dtype']
+        )
+        clip_processor = CLIPProcessor.from_pretrained("wkcn/TinyCLIP-ViT-40M-32-Text-19M-LAION400M")
+    clip_model = clip_model.to(params['device'], dtype=params['weight_dtype'])
+    clip_model.requires_grad_(False)
+
+
+    pipe = FluxPipeline(noise_scheduler,
+                    vae,
+                    text_encoder_one,
+                    tokenizer_one,
+                    text_encoder_two,
+                    tokenizer_two,
+                    transformer,
+                   )
+    pipe.set_progress_bar_config(disable=True)
+
+    return {
+        'transformer': transformer,
+        'vae': vae,
+        'clip_model': clip_model,
+        'clip_processor': clip_processor,
+        'tokenizers': [tokenizer_one, tokenizer_two],
+        'text_encoders': [text_encoder_one,text_encoder_two],
+        'noise_scheduler': noise_scheduler
+    }, pipe
+
+def save_checkpoint(networks, save_path, weight_dtype):
+    """
+    Save network weights and perform cleanup
+    
+    Args:
+        networks: Dictionary of LoRA networks to save
+        save_path: Path to save the checkpoints
+        weight_dtype: Data type for the weights
+    """
+    print("Saving checkpoint...")
+    
+    try:
+        # Create save directory if it doesn't exist
+        os.makedirs(save_path, exist_ok=True)
+        
+        # Save each network's weights
+        for net_idx, network in networks.items():
+            save_name = f"{save_path}/slider_{net_idx}.pt"
+            try:
+                network.save_weights(
+                    save_name,
+                    dtype=weight_dtype,
+                )
+            except Exception as e:
+                print(f"Error saving network {net_idx}: {str(e)}")
+                continue
+                
+        # Cleanup
+        torch.cuda.empty_cache()
+        gc.collect()
+        
+        print("Checkpoint saved successfully.")
+        
+    except Exception as e:
+        print(f"Error during checkpoint saving: {str(e)}")
+        
+    finally:
+        # Ensure memory is cleaned up even if save fails
+        torch.cuda.empty_cache()
+        gc.collect()

utils/prompt_util.py

@@ -0,0 +1,80 @@
+import anthropic
+client = anthropic.Anthropic()
+from typing import List, Optional
+
+def claude_generate_prompts_sliders(prompt, 
+                             num_prompts=20,
+                             temperature=0.2, 
+                             max_tokens=2000, 
+                             frequency_penalty=0.0,
+                             model="claude-3-5-sonnet-20240620",
+                             verbose=False):
+    assistant_prompt =  f''' You are an expert in writing diverse image captions. When i provide a prompt, I want you to give me {num_prompts} alternative prompts that is similar to the provided prompt but produces diverse images. Be creative and make sure the original subjects in the original prompt are present in your prompts. Make sure that you end the prompts with keywords that will produce high quality images like ",detailed, 8k" or ",hyper-realistic, 4k".
+
+Give me the expanded prompts in the style of a list. start with a [ and end with ] do not add any special characters like \n 
+I need you to give me only the python list and nothing else. Do not explain yourself
+
+example output format:
+["prompt1", "prompt2", ...]
+'''
+    
+    user_prompt = prompt
+    
+    message=[
+        {
+            "role": "user", 
+            "content": [
+                {
+                    "type": "text",
+                    "text": user_prompt
+                }
+            ]
+        }
+            ]
+    
+    output = client.messages.create(
+        model=model,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        system=assistant_prompt,
+        messages=message
+    )
+    content = output.content[0].text
+    return content
+
+
+def expand_prompts(concept_prompts: List[str], diverse_prompt_num: int, args) -> List[str]:
+    """
+    Expand the input prompts using Claude if requested.
+    
+    Args:
+        concept_prompts: Initial list of prompts
+        diverse_prompt_num: Number of variations to generate per prompt
+        args: Training arguments
+        
+    Returns:
+        List of expanded prompts
+    """
+    diverse_prompts = []
+    
+    if diverse_prompt_num != 0:
+        for prompt in concept_prompts:
+            try:
+                claude_generated_prompts = claude_generate_prompts_sliders(
+                    prompt=prompt,
+                    num_prompts=diverse_prompt_num,
+                    temperature=0.2,
+                    max_tokens=8000,
+                    frequency_penalty=0.0, 
+                    model="claude-3-5-sonnet-20240620",
+                    verbose=False
+                )
+                diverse_prompts.extend(eval(claude_generated_prompts))
+            except Exception as e:
+                print(f"Error with Claude response: {e}")
+                diverse_prompts.append(prompt)
+    else:
+        diverse_prompts = concept_prompts
+        
+    print(f"Using prompts: {diverse_prompts}")
+    return diverse_prompts

utils/train_util.py

@@ -0,0 +1,722 @@
+from typing import Optional, Union
+
+import torch
+
+from transformers import CLIPTextModel, CLIPTokenizer
+from diffusers import UNet2DConditionModel, SchedulerMixin, FluxImg2ImgPipeline
+from diffusers.image_processor import VaeImageProcessor
+# from model_util import SDXL_TEXT_ENCODER_TYPE
+from diffusers.utils.torch_utils import randn_tensor
+
+from tqdm import tqdm
+
+UNET_IN_CHANNELS = 4  # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
+VAE_SCALE_FACTOR = 8  # 2 ** (len(vae.config.block_out_channels) - 1) = 8
+
+UNET_ATTENTION_TIME_EMBED_DIM = 256  # XL
+TEXT_ENCODER_2_PROJECTION_DIM = 1280
+UNET_PROJECTION_CLASS_EMBEDDING_INPUT_DIM = 2816
+
+
+def get_random_noise(
+    batch_size: int, height: int, width: int, generator: torch.Generator = None
+) -> torch.Tensor:
+    return torch.randn(
+        (
+            batch_size,
+            UNET_IN_CHANNELS,
+            height // VAE_SCALE_FACTOR,  # 縦と横これであってるのかわからないけど、どっちにしろ大きな問題は発生しないのでこれでいいや
+            width // VAE_SCALE_FACTOR,
+        ),
+        generator=generator,
+        device="cpu",
+    )
+
+
+# https://www.crosslabs.org/blog/diffusion-with-offset-noise
+def apply_noise_offset(latents: torch.FloatTensor, noise_offset: float):
+    latents = latents + noise_offset * torch.randn(
+        (latents.shape[0], latents.shape[1], 1, 1), device=latents.device
+    )
+    return latents
+
+
+def get_initial_latents(
+    scheduler: SchedulerMixin,
+    n_imgs: int,
+    height: int,
+    width: int,
+    n_prompts: int,
+    generator=None,
+) -> torch.Tensor:
+    noise = get_random_noise(n_imgs, height, width, generator=generator).repeat(
+        n_prompts, 1, 1, 1
+    )
+
+    latents = noise * scheduler.init_noise_sigma
+
+    return latents
+
+
+def text_tokenize(
+    tokenizer: CLIPTokenizer,  # 普通ならひとつ、XLならふたつ！
+    prompts: list[str],
+):
+    return tokenizer(
+        prompts,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+        return_tensors="pt",
+    ).input_ids
+
+
+def text_encode(text_encoder: CLIPTextModel, tokens):
+    return text_encoder(tokens.to(text_encoder.device))[0]
+
+
+def encode_prompts(
+    tokenizer: CLIPTokenizer,
+    text_encoder: CLIPTokenizer,
+    prompts: list[str],
+):
+
+    text_tokens = text_tokenize(tokenizer, prompts)
+    text_embeddings = text_encode(text_encoder, text_tokens)
+    
+    
+
+    return text_embeddings
+
+
+# https://github.com/huggingface/diffusers/blob/78922ed7c7e66c20aa95159c7b7a6057ba7d590d/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L334-L348
+def text_encode_xl(
+    text_encoder,
+    tokens: torch.FloatTensor,
+    num_images_per_prompt: int = 1,
+):
+    prompt_embeds = text_encoder(
+        tokens.to(text_encoder.device), output_hidden_states=True
+    )
+    pooled_prompt_embeds = prompt_embeds[0]
+    prompt_embeds = prompt_embeds.hidden_states[-2]  # always penultimate layer
+
+    bs_embed, seq_len, _ = prompt_embeds.shape
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+    prompt_embeds = prompt_embeds.view(bs_embed * num_images_per_prompt, seq_len, -1)
+
+    return prompt_embeds, pooled_prompt_embeds
+
+
+def encode_prompts_xl(
+    tokenizers,
+    text_encoders,
+    prompts: list[str],
+    num_images_per_prompt: int = 1,
+) -> tuple[torch.FloatTensor, torch.FloatTensor]:
+    # text_encoder and text_encoder_2's penuultimate layer's output
+    text_embeds_list = []
+    pooled_text_embeds = None  # always text_encoder_2's pool
+
+    for tokenizer, text_encoder in zip(tokenizers, text_encoders):
+        text_tokens_input_ids = text_tokenize(tokenizer, prompts)
+        text_embeds, pooled_text_embeds = text_encode_xl(
+            text_encoder, text_tokens_input_ids, num_images_per_prompt
+        )
+
+        text_embeds_list.append(text_embeds)
+
+    bs_embed = pooled_text_embeds.shape[0]
+    pooled_text_embeds = pooled_text_embeds.repeat(1, num_images_per_prompt).view(
+        bs_embed * num_images_per_prompt, -1
+    )
+
+    return torch.concat(text_embeds_list, dim=-1), pooled_text_embeds
+
+
+def concat_embeddings(
+    unconditional: torch.FloatTensor,
+    conditional: torch.FloatTensor,
+    n_imgs: int,
+):
+    return torch.cat([unconditional, conditional]).repeat_interleave(n_imgs, dim=0)
+
+
+# ref: https://github.com/huggingface/diffusers/blob/0bab447670f47c28df60fbd2f6a0f833f75a16f5/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L721
+def predict_noise(
+    unet: UNet2DConditionModel,
+    scheduler: SchedulerMixin,
+    timestep: int,  # 現在のタイムステップ
+    latents: torch.FloatTensor,
+    text_embeddings: torch.FloatTensor,  # uncond な text embed と cond な text embed を結合したもの
+    guidance_scale=7.5,
+) -> torch.FloatTensor:
+    latent_model_input = latents
+    if guidance_scale!=0:
+        # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+        latent_model_input = torch.cat([latents] * 2)
+
+    latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)
+
+    # predict the noise residual
+    noise_pred = unet(
+        latent_model_input,
+        timestep,
+        encoder_hidden_states=text_embeddings,
+    ).sample
+
+    # perform guidance
+    if guidance_scale != 1 and guidance_scale!=0:
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+    return noise_pred
+
+
+# ref: https://github.com/huggingface/diffusers/blob/0bab447670f47c28df60fbd2f6a0f833f75a16f5/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L746
+@torch.no_grad()
+def diffusion(
+    unet: UNet2DConditionModel,
+    scheduler: SchedulerMixin,
+    latents: torch.FloatTensor,  # ただのノイズだけのlatents
+    text_embeddings: torch.FloatTensor,
+    total_timesteps: int = 1000,
+    start_timesteps=0,
+    guidance_scale=1,
+    composition=False,
+    **kwargs,
+):
+    # latents_steps = []
+
+    for timestep in scheduler.timesteps[start_timesteps:total_timesteps]:
+        if not composition:
+            noise_pred = predict_noise(
+                unet, scheduler, timestep, latents, text_embeddings, guidance_scale=guidance_scale
+            )
+            if guidance_scale==1:
+                _, noise_pred = noise_pred.chunk(2)
+        else:
+            for idx in range(text_embeddings.shape[0]):
+                pred = predict_noise(
+                    unet, scheduler, timestep, latents, text_embeddings[idx:idx+1], guidance_scale=1
+                )
+                uncond, pred = noise_pred.chunk(2)
+                if idx == 0:
+                    noise_pred = guidance_scale * pred
+                else:
+                    noise_pred += guidance_scale * pred
+            noise_pred += uncond
+        
+        
+        # compute the previous noisy sample x_t -> x_t-1
+        latents = scheduler.step(noise_pred, timestep, latents).prev_sample
+
+    # return latents_steps
+    return latents
+
+
+def rescale_noise_cfg(
+    noise_cfg: torch.FloatTensor, noise_pred_text, guidance_rescale=0.0
+):
+    """
+    Rescale `noise_cfg` according to `guidance_rescale`. Based on findings of [Common Diffusion Noise Schedules and
+    Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf). See Section 3.4
+    """
+    std_text = noise_pred_text.std(
+        dim=list(range(1, noise_pred_text.ndim)), keepdim=True
+    )
+    std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+    # rescale the results from guidance (fixes overexposure)
+    noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+    # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images
+    noise_cfg = (
+        guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+    )
+
+    return noise_cfg
+
+
+def predict_noise_xl(
+    unet: UNet2DConditionModel,
+    scheduler: SchedulerMixin,
+    timestep: int,  # 現在のタイムステップ
+    latents: torch.FloatTensor,
+    text_embeddings: torch.FloatTensor,  # uncond な text embed と cond な text embed を結合したもの
+    add_text_embeddings: torch.FloatTensor,  # pooled なやつ
+    add_time_ids: torch.FloatTensor,
+    guidance_scale=7.5,
+    guidance_rescale=0.7,
+) -> torch.FloatTensor:
+    # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
+    latent_model_input = latents
+    if guidance_scale !=0:
+        latent_model_input = torch.cat([latents] * 2)
+    
+    latent_model_input = scheduler.scale_model_input(latent_model_input, timestep)
+
+    added_cond_kwargs = {
+        "text_embeds": add_text_embeddings,
+        "time_ids": add_time_ids,
+    }
+
+    # predict the noise residual
+    noise_pred = unet(
+        latent_model_input,
+        timestep,
+        encoder_hidden_states=text_embeddings,
+        added_cond_kwargs=added_cond_kwargs,
+    ).sample
+    # perform guidance
+    if guidance_scale != 1 and guidance_scale!=0:
+        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+        noise_pred = noise_pred_uncond + guidance_scale * (
+            noise_pred_text - noise_pred_uncond
+        )
+
+    return noise_pred
+    # # perform guidance
+    # noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+    # guided_target = noise_pred_uncond + guidance_scale * (
+    #     noise_pred_text - noise_pred_uncond
+    # )
+
+    # # https://github.com/huggingface/diffusers/blob/7a91ea6c2b53f94da930a61ed571364022b21044/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L775
+    # noise_pred = rescale_noise_cfg(
+    #     noise_pred, noise_pred_text, guidance_rescale=guidance_rescale
+    # )
+
+    # return guided_target
+
+
+@torch.no_grad()
+def diffusion_xl(
+    unet: UNet2DConditionModel,
+    scheduler: SchedulerMixin,
+    latents: torch.FloatTensor,  # ただのノイズだけのlatents
+    text_embeddings: tuple[torch.FloatTensor, torch.FloatTensor],
+    add_text_embeddings: torch.FloatTensor,  # pooled なやつ
+    add_time_ids: torch.FloatTensor,
+    guidance_scale: float = 1.0,
+    total_timesteps: int = 1000,
+    start_timesteps=0,
+    composition=False,
+):
+    # latents_steps = []
+
+    for timestep in scheduler.timesteps[start_timesteps:total_timesteps]:
+        if not composition:
+            noise_pred = predict_noise_xl(
+                unet,
+                scheduler,
+                timestep,
+                latents,
+                text_embeddings,
+                add_text_embeddings,
+                add_time_ids,
+                guidance_scale=guidance_scale,
+                guidance_rescale=0.7,
+            )
+            if guidance_scale==1:
+                _, noise_pred = noise_pred.chunk(2)
+        # compute the previous noisy sample x_t -> x_t-1
+        latents = scheduler.step(noise_pred, timestep, latents).prev_sample
+
+    # return latents_steps
+    return latents
+
+
+# for XL
+def get_add_time_ids(
+    height: int,
+    width: int,
+    dynamic_crops: bool = False,
+    dtype: torch.dtype = torch.float32,
+):
+    if dynamic_crops:
+        # random float scale between 1 and 3
+        random_scale = torch.rand(1).item() * 2 + 1
+        original_size = (int(height * random_scale), int(width * random_scale))
+        # random position
+        crops_coords_top_left = (
+            torch.randint(0, original_size[0] - height, (1,)).item(),
+            torch.randint(0, original_size[1] - width, (1,)).item(),
+        )
+        target_size = (height, width)
+    else:
+        original_size = (height, width)
+        crops_coords_top_left = (0, 0)
+        target_size = (height, width)
+
+    # this is expected as 6
+    add_time_ids = list(original_size + crops_coords_top_left + target_size)
+
+    # this is expected as 2816
+    passed_add_embed_dim = (
+        UNET_ATTENTION_TIME_EMBED_DIM * len(add_time_ids)  # 256 * 6
+        + TEXT_ENCODER_2_PROJECTION_DIM  # + 1280
+    )
+    if passed_add_embed_dim != UNET_PROJECTION_CLASS_EMBEDDING_INPUT_DIM:
+        raise ValueError(
+            f"Model expects an added time embedding vector of length {UNET_PROJECTION_CLASS_EMBEDDING_INPUT_DIM}, but a vector of {passed_add_embed_dim} was created. The model has an incorrect config. Please check `unet.config.time_embedding_type` and `text_encoder_2.config.projection_dim`."
+        )
+
+    add_time_ids = torch.tensor([add_time_ids], dtype=dtype)
+    return add_time_ids
+
+
+def get_optimizer(name: str):
+    name = name.lower()
+
+    if name.startswith("dadapt"):
+        import dadaptation
+
+        if name == "dadaptadam":
+            return dadaptation.DAdaptAdam
+        elif name == "dadaptlion":
+            return dadaptation.DAdaptLion
+        else:
+            raise ValueError("DAdapt optimizer must be dadaptadam or dadaptlion")
+
+    elif name.endswith("8bit"):  # 検証してない
+        import bitsandbytes as bnb
+
+        if name == "adam8bit":
+            return bnb.optim.Adam8bit
+        elif name == "lion8bit":
+            return bnb.optim.Lion8bit
+        else:
+            raise ValueError("8bit optimizer must be adam8bit or lion8bit")
+
+    else:
+        if name == "adam":
+            return torch.optim.Adam
+        elif name == "adamw":
+            return torch.optim.AdamW
+        elif name == "lion":
+            from lion_pytorch import Lion
+
+            return Lion
+        elif name == "prodigy":
+            import prodigyopt
+            
+            return prodigyopt.Prodigy
+        else:
+            raise ValueError("Optimizer must be adam, adamw, lion or Prodigy")
+
+@torch.no_grad()
+def get_noisy_image(
+    image,
+    vae,
+    unet,
+    scheduler,
+    timesteps_to = 1000,
+    generator=None,
+    **kwargs,
+):
+    # latents_steps = []
+    vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
+
+    device = vae.device
+    image = image_processor.preprocess(image).to(device).to(vae.dtype)
+
+    init_latents = vae.encode(image).latents
+    
+    init_latents = vae.config.scaling_factor * init_latents
+
+    init_latents = torch.cat([init_latents], dim=0)
+
+    shape = init_latents.shape
+
+    noise = randn_tensor(shape, generator=generator, device=device)
+
+    timestep = scheduler.timesteps[timesteps_to:timesteps_to+1]
+    # get latents
+    init_latents = scheduler.add_noise(init_latents, noise, timestep)
+    
+    return init_latents, noise
+
+
+def get_lr_scheduler(
+    name: Optional[str],
+    optimizer: torch.optim.Optimizer,
+    max_iterations: Optional[int],
+    lr_min: Optional[float],
+    **kwargs,
+):
+    if name == "cosine":
+        return torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=max_iterations, eta_min=lr_min, **kwargs
+        )
+    elif name == "cosine_with_restarts":
+        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer, T_0=max_iterations // 10, T_mult=2, eta_min=lr_min, **kwargs
+        )
+    elif name == "step":
+        return torch.optim.lr_scheduler.StepLR(
+            optimizer, step_size=max_iterations // 100, gamma=0.999, **kwargs
+        )
+    elif name == "constant":
+        return torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1, **kwargs)
+    elif name == "linear":
+        return torch.optim.lr_scheduler.LinearLR(
+            optimizer, factor=0.5, total_iters=max_iterations // 100, **kwargs
+        )
+    else:
+        raise ValueError(
+            "Scheduler must be cosine, cosine_with_restarts, step, linear or constant"
+        )
+
+
+def get_random_resolution_in_bucket(bucket_resolution: int = 512) -> tuple[int, int]:
+    max_resolution = bucket_resolution
+    min_resolution = bucket_resolution // 2
+
+    step = 64
+
+    min_step = min_resolution // step
+    max_step = max_resolution // step
+
+    height = torch.randint(min_step, max_step, (1,)).item() * step
+    width = torch.randint(min_step, max_step, (1,)).item() * step
+
+    return height, width
+
+
+
+def _get_t5_prompt_embeds(
+    text_encoder,
+    tokenizer,
+    prompt,
+    max_sequence_length=512,
+    device=None,
+    dtype=None
+):
+    """Helper function to get T5 embeddings in Flux format"""
+    device = device or text_encoder.device
+    dtype = dtype or text_encoder.dtype
+
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=max_sequence_length,
+        truncation=True,
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_tensors="pt",
+    )
+    text_input_ids = text_inputs.input_ids
+
+    prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=False)[0]
+    prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+    return prompt_embeds
+
+def _get_clip_prompt_embeds(
+    text_encoder,
+    tokenizer,
+    prompt,
+    device=None,
+):
+    """Helper function to get CLIP embeddings in Flux format"""
+    device = device or text_encoder.device
+
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    text_inputs = tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+        return_overflowing_tokens=False,
+        return_length=False,
+        return_tensors="pt",
+    )
+
+    text_input_ids = text_inputs.input_ids
+    prompt_embeds = text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+    # Use pooled output for Flux
+    prompt_embeds = prompt_embeds.pooler_output
+    prompt_embeds = prompt_embeds.to(dtype=text_encoder.dtype, device=device)
+
+    return prompt_embeds
+
+
+
+
+
+@torch.no_grad()
+def get_noisy_image_flux(
+    image,
+    vae,
+    transformer,
+    scheduler,
+    timesteps_to=1000,
+    generator=None,
+    params = None
+):
+    """
+    Gets noisy latents for a given image using Flux pipeline approach.
+    
+    Args:
+        image (Union[PIL.Image.Image, torch.Tensor]): Input image
+        vae (AutoencoderKL): Flux VAE model
+        transformer (FluxTransformer2DModel): Flux transformer model
+        scheduler (FlowMatchEulerDiscreteScheduler): Flux noise scheduler
+        timesteps_to (int, optional): Target timestep. Defaults to 1000.
+        generator (torch.Generator, optional): Random generator for reproducibility.
+            
+    Returns:
+        tuple: (noisy_latents, noise) - Both in packed Flux format
+    """
+
+    vae_scale_factor = params['vae_scale_factor']
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor * 2)
+    
+    image = image_processor.preprocess(image, height=params['height'], width=params['width'])
+    image = image.to(dtype=torch.float32)
+    
+    # 5. Prepare latent variables
+    num_channels_latents = transformer.config.in_channels // 4
+
+    latents, latent_image_ids = prepare_latents_flux(
+        image,
+        timesteps_to.repeat(params['batchsize']),
+        params['batchsize'],
+        num_channels_latents,
+        params['height'],
+        params['width'],
+        transformer.dtype,
+        transformer.device,
+        generator,
+        None,
+        vae_scale_factor,
+        vae,
+        scheduler
+    )
+    
+    return latents, latent_image_ids
+
+
+def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+    """
+    Pack latents into Flux's 2x2 patch format
+    """
+    latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+    latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+    return latents
+
+
+def _unpack_latents(latents, height, width, vae_scale_factor):
+    """
+    Unpack latents from Flux's 2x2 patch format back to image space
+    """
+    batch_size, num_patches, channels = latents.shape
+
+    # Account for VAE compression and packing
+    height = 2 * (int(height) // (vae_scale_factor * 2))
+    width = 2 * (int(width) // (vae_scale_factor * 2))
+
+    latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+    latents = latents.reshape(batch_size, channels // (2 * 2), height, width)
+
+    return latents
+
+def _prepare_latent_image_ids(batch_size, height, width, device, dtype):
+    latent_image_ids = torch.zeros(height, width, 3)
+    latent_image_ids[..., 1] = latent_image_ids[..., 1] + torch.arange(height)[:, None]
+    latent_image_ids[..., 2] = latent_image_ids[..., 2] + torch.arange(width)[None, :]
+
+    latent_image_id_height, latent_image_id_width, latent_image_id_channels = latent_image_ids.shape
+
+    latent_image_ids = latent_image_ids.reshape(
+        latent_image_id_height * latent_image_id_width, latent_image_id_channels
+    )
+
+    return latent_image_ids.to(device=device, dtype=dtype)
+
+
+def prepare_latents_flux(
+    image,
+    timestep,
+    batch_size,
+    num_channels_latents,
+    height,
+    width,
+    dtype,
+    device,
+    generator,
+    latents=None,
+    vae_scale_factor=None,
+    vae=None,
+    scheduler=None
+):
+    if isinstance(generator, list) and len(generator) != batch_size:
+        raise ValueError(
+            f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+            f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+        )
+
+    # VAE applies 8x compression on images but we must also account for packing which requires
+    # latent height and width to be divisible by 2.
+    height = 2 * (int(height) // (vae_scale_factor * 2))
+    width = 2 * (int(width) // (vae_scale_factor * 2))
+    shape = (batch_size, num_channels_latents, height, width)
+    latent_image_ids = _prepare_latent_image_ids(batch_size, height // 2, width // 2, device, dtype)
+
+    if latents is not None:
+        return latents.to(device=device, dtype=dtype), latent_image_ids
+
+    image = image.to(device=device, dtype=dtype)
+    image_latents = _encode_vae_image(vae=vae, image=image, generator=generator)
+    if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+        # expand init_latents for batch_size
+        additional_image_per_prompt = batch_size // image_latents.shape[0]
+        image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+    elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+        raise ValueError(
+            f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+        )
+    else:
+        image_latents = torch.cat([image_latents], dim=0)
+
+    noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+    latents = scheduler.scale_noise(image_latents, timestep, noise)
+    latents = _pack_latents(latents, batch_size, num_channels_latents, height, width)
+    return latents, latent_image_ids
+
+
+def _encode_vae_image(vae, image: torch.Tensor, generator: torch.Generator):
+    if isinstance(generator, list):
+        image_latents = [
+            retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i])
+            for i in range(image.shape[0])
+        ]
+        image_latents = torch.cat(image_latents, dim=0)
+    else:
+        image_latents = retrieve_latents(vae.encode(image), generator=generator)
+
+    image_latents = (image_latents - vae.config.shift_factor) * vae.config.scaling_factor
+    return image_latents
+
+
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")

utils/utils.py

@@ -0,0 +1,945 @@
+import anthropic
+client = anthropic.Anthropic()
+from diffusers.image_processor import VaeImageProcessor
+from typing import List, Optional
+import argparse
+import ast
+import pandas as pd
+from pathlib import Path
+from diffusers import DiffusionPipeline, UNet2DConditionModel, LCMScheduler, AutoencoderTiny
+from huggingface_hub import hf_hub_download
+import gc
+import torch.nn.functional as F
+import os
+import torch
+from tqdm.auto import tqdm
+import time, datetime
+import numpy as np
+from torch.optim import AdamW
+from contextlib import ExitStack
+from safetensors.torch import load_file
+import torch.nn as nn
+import random
+from transformers import CLIPModel
+
+import sys
+import argparse
+import wandb
+from diffusers import AutoencoderKL
+from diffusers.image_processor import VaeImageProcessor
+
+sys.path.append('../')
+from utils.lora import LoRANetwork, DEFAULT_TARGET_REPLACE, UNET_TARGET_REPLACE_MODULE_CONV
+
+from transformers import logging
+logging.set_verbosity_warning()
+import matplotlib.pyplot as plt
+from diffusers import logging
+logging.set_verbosity_error()
+modules = DEFAULT_TARGET_REPLACE
+modules += UNET_TARGET_REPLACE_MODULE_CONV
+import torch
+import torch.nn.functional as F
+from sklearn.decomposition import PCA
+import random
+import gc
+import diffusers
+from diffusers import DiffusionPipeline, FluxPipeline
+from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel, LMSDiscreteScheduler, SchedulerMixin
+from diffusers.loaders import AttnProcsLayers
+from diffusers.models.attention_processor import LoRAAttnProcessor, AttentionProcessor
+from typing import Any, Dict, List, Optional, Tuple, Union
+from diffusers.pipelines.stable_diffusion_xl import StableDiffusionXLPipelineOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+import inspect
+import os
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from diffusers.pipelines import StableDiffusionXLPipeline
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import retrieve_timesteps
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import XLA_AVAILABLE
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
+
+from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusionPipelineOutput
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+
+import sys
+sys.path.append('../.')
+from utils.flux_utils import *
+import random
+
+import torch
+from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.16,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+def claude_generate_prompts_sliders(prompt, 
+                             num_prompts=20,
+                             temperature=0.2, 
+                             max_tokens=2000, 
+                             frequency_penalty=0.0,
+                             model="claude-3-5-sonnet-20240620",
+                             verbose=False,
+                             train_type='concept'):
+    gpt_assistant_prompt =  f''' You are an expert in writing diverse image captions. When i provide a prompt, I want you to give me {num_prompts} alternative prompts that is similar to the provided prompt but produces diverse images. Be creative and make sure the original subjects in the original prompt are present in your prompts. Make sure that you end the prompts with keywords that will produce high quality images like ",detailed, 8k" or ",hyper-realistic, 4k".
+
+Give me the expanded prompts in the style of a list. start with a [ and end with ] do not add any special characters like \n 
+I need you to give me only the python list and nothing else. Do not explain yourself
+
+example output format:
+["prompt1", "prompt2", ...]
+'''
+    
+    if train_type == 'art':
+        gpt_assistant_prompt =  f'''You are an expert in writing art image captions. I want you to generate prompts that would create diverse artwork images. 
+    Your role is to give me {num_prompts} diverse prompts that will make the image-generation model to output creative and interesting artwork images with unique and diverse artistic styles. A prompt could like "an <object/landscape> in the style of <an artist>" or "an <object/landscape> in the style of <an artistic style (e.g. cubism)>". make sure that you end the prompts with enhancing keywords like ",detailed, 8k" or ",hyper-realistic, 4k". 
+    
+   Give me the prompts in the style of a list. start with a [ and end with ] do not add any special characters like \n 
+I need you to give me only the python list and nothing else. Do not explain yourself
+
+example output format:
+["prompt1", "prompt2", ...]
+    '''
+    # if 'dog' in prompt:
+    #     gpt_assistant_prompt =  f'''You are an expert in prompting text-image generation models. I want you to generate simple prompts that would trigger the image generation model to generate a unique dog breeds. 
+    # Your role is to give me {num_prompts} diverse prompts that will make the image-generation model to output diverse and interesting dog breeds with unique and diverse looks. make sure that you end the prompts with enhancing keywords like ",detailed, 8k" or ",hyper-realistic, 4k". 
+    
+    # Be creative and make sure to remember diversity is the key. Give me the prompts in the form of a list. start with a [ and end with ] do not add any special characters like \n 
+    # I need you to give me only the python list and nothing else. Do not explain yourself
+
+    # example output format:
+    # ["prompt1", "prompt2", ...]
+    # '''        
+
+    if train_type == 'artclaudesemantics':
+        gpt_assistant_prompt =  f'''You are an expert in prompting text-image generation models. I want you to generate simple prompts that would trigger the image generation model to generate a unique artistic images but DO NOT SPECIFY THE ART STYLE. 
+    Your role is to give me {num_prompts} diverse prompts that will make the image-generation model to output diverse and interesting art images. Usually like "<some object or scene> in the style of " or "<some object or scene> in style of". Always end your prompts with "in the style of" so that i can manually add the style i want. make sure that you end the prompts with enhancing keywords like ",detailed, 8k" or ",hyper-realistic, 4k". 
+    
+    Be creative and make sure to remember diversity is the key. Give me the prompts in the form of a list. start with a [ and end with ] do not add any special characters like \n 
+    I need you to give me only the python list and nothing else. Do not explain yourself
+
+    example output format:
+    ["prompt1", "prompt2", ...]
+    '''
+    gpt_user_prompt = prompt
+    gpt_prompt = gpt_assistant_prompt, gpt_user_prompt
+    message=[
+        {
+            "role": "user", 
+            "content": [
+                {
+                    "type": "text",
+                    "text": gpt_user_prompt
+                }
+            ]
+        }
+            ]
+    
+    output = client.messages.create(
+        model=model,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        system=gpt_assistant_prompt,
+        messages=message
+    )
+    content = output.content[0].text
+    return content
+
+def normalize_image(image):
+    mean = torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1).to(image.device)
+    std = torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1).to(image.device)
+    return (image - mean) / std
+
+
+@torch.no_grad()
+def call_sdxl(
+    self,
+    prompt: Union[str, List[str]] = None,
+    prompt_2: Optional[Union[str, List[str]]] = None,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_inference_steps: int = 50,
+    timesteps: List[int] = None,
+    sigmas: List[float] = None,
+    denoising_end: Optional[float] = None,
+    guidance_scale: float = 5.0,
+    negative_prompt: Optional[Union[str, List[str]]] = None,
+    negative_prompt_2: Optional[Union[str, List[str]]] = None,
+    num_images_per_prompt: Optional[int] = 1,
+    eta: float = 0.0,
+    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    latents: Optional[torch.Tensor] = None,
+    prompt_embeds: Optional[torch.Tensor] = None,
+    negative_prompt_embeds: Optional[torch.Tensor] = None,
+    pooled_prompt_embeds: Optional[torch.Tensor] = None,
+    negative_pooled_prompt_embeds: Optional[torch.Tensor] = None,
+    ip_adapter_image: Optional[PipelineImageInput] = None,
+    ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+    output_type: Optional[str] = "pil",
+    return_dict: bool = True,
+    cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+    guidance_rescale: float = 0.0,
+    original_size: Optional[Tuple[int, int]] = None,
+    crops_coords_top_left: Tuple[int, int] = (0, 0),
+    target_size: Optional[Tuple[int, int]] = None,
+    negative_original_size: Optional[Tuple[int, int]] = None,
+    negative_crops_coords_top_left: Tuple[int, int] = (0, 0),
+    negative_target_size: Optional[Tuple[int, int]] = None,
+    clip_skip: Optional[int] = None,
+    callback_on_step_end: Optional[
+        Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+    ] = None,
+    callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+    save_timesteps = None,
+    clip=None,
+    use_clip=True,
+    encoder='clip',
+):
+
+    callback = None
+    callback_steps = None
+
+    if callback is not None:
+        deprecate(
+            "callback",
+            "1.0.0",
+            "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+        )
+    if callback_steps is not None:
+        deprecate(
+            "callback_steps",
+            "1.0.0",
+            "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+        )
+
+    if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+        callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+    # 0. Default height and width to unet
+    height = height or self.default_sample_size * self.vae_scale_factor
+    width = width or self.default_sample_size * self.vae_scale_factor
+
+    original_size = original_size or (height, width)
+    target_size = target_size or (height, width)
+
+    # 1. Check inputs. Raise error if not correct
+    self.check_inputs(
+        prompt,
+        prompt_2,
+        height,
+        width,
+        callback_steps,
+        negative_prompt,
+        negative_prompt_2,
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+        ip_adapter_image,
+        ip_adapter_image_embeds,
+        callback_on_step_end_tensor_inputs,
+    )
+
+    self._guidance_scale = guidance_scale
+    self._guidance_rescale = guidance_rescale
+    self._clip_skip = clip_skip
+    self._cross_attention_kwargs = cross_attention_kwargs
+    self._denoising_end = denoising_end
+    self._interrupt = False
+
+    # 2. Define call parameters
+    if prompt is not None and isinstance(prompt, str):
+        batch_size = 1
+    elif prompt is not None and isinstance(prompt, list):
+        batch_size = len(prompt)
+    else:
+        batch_size = prompt_embeds.shape[0]
+
+    device = self._execution_device
+
+    # 3. Encode input prompt
+    lora_scale = (
+        self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
+    )
+
+    (
+        prompt_embeds,
+        negative_prompt_embeds,
+        pooled_prompt_embeds,
+        negative_pooled_prompt_embeds,
+    ) = self.encode_prompt(
+        prompt=prompt,
+        prompt_2=prompt_2,
+        device=device,
+        num_images_per_prompt=num_images_per_prompt,
+        do_classifier_free_guidance=self.do_classifier_free_guidance,
+        negative_prompt=negative_prompt,
+        negative_prompt_2=negative_prompt_2,
+        prompt_embeds=prompt_embeds,
+        negative_prompt_embeds=negative_prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+        lora_scale=lora_scale,
+        clip_skip=self.clip_skip,
+    )
+
+    # 4. Prepare timesteps
+    timesteps, num_inference_steps = retrieve_timesteps(
+        self.scheduler, num_inference_steps, device, timesteps, sigmas
+    )
+
+    # 5. Prepare latent variables
+    num_channels_latents = self.unet.config.in_channels
+    latents = self.prepare_latents(
+        batch_size * num_images_per_prompt,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+
+    # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+    extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+    # 7. Prepare added time ids & embeddings
+    add_text_embeds = pooled_prompt_embeds
+    if self.text_encoder_2 is None:
+        text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+    else:
+        text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+    add_time_ids = self._get_add_time_ids(
+        original_size,
+        crops_coords_top_left,
+        target_size,
+        dtype=prompt_embeds.dtype,
+        text_encoder_projection_dim=text_encoder_projection_dim,
+    )
+    if negative_original_size is not None and negative_target_size is not None:
+        negative_add_time_ids = self._get_add_time_ids(
+            negative_original_size,
+            negative_crops_coords_top_left,
+            negative_target_size,
+            dtype=prompt_embeds.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+    else:
+        negative_add_time_ids = add_time_ids
+
+    if self.do_classifier_free_guidance:
+        prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+        add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0)
+        add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+    prompt_embeds = prompt_embeds.to(device)
+    add_text_embeds = add_text_embeds.to(device)
+    add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1)
+
+    if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+        image_embeds = self.prepare_ip_adapter_image_embeds(
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            device,
+            batch_size * num_images_per_prompt,
+            self.do_classifier_free_guidance,
+        )
+
+    # 8. Denoising loop
+    num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+
+    # 8.1 Apply denoising_end
+    if (
+        self.denoising_end is not None
+        and isinstance(self.denoising_end, float)
+        and self.denoising_end > 0
+        and self.denoising_end < 1
+    ):
+        discrete_timestep_cutoff = int(
+            round(
+                self.scheduler.config.num_train_timesteps
+                - (self.denoising_end * self.scheduler.config.num_train_timesteps)
+            )
+        )
+        num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps)))
+        timesteps = timesteps[:num_inference_steps]
+
+    # 9. Optionally get Guidance Scale Embedding
+    timestep_cond = None
+    if self.unet.config.time_cond_proj_dim is not None:
+        guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
+        timestep_cond = self.get_guidance_scale_embedding(
+            guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+        ).to(device=device, dtype=latents.dtype)
+
+    self._num_timesteps = len(timesteps)
+    clip_features = []
+    # with self.progress_bar(total=num_inference_steps) as progress_bar:
+    for i, t in enumerate(timesteps):
+        if self.interrupt:
+            continue
+
+        # expand the latents if we are doing classifier free guidance
+        latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+
+        latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+        # predict the noise residual
+        added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids}
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            added_cond_kwargs["image_embeds"] = image_embeds
+        noise_pred = self.unet(
+            latent_model_input,
+            t,
+            encoder_hidden_states=prompt_embeds,
+            timestep_cond=timestep_cond,
+            cross_attention_kwargs=self.cross_attention_kwargs,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )[0]
+
+        # perform guidance
+        if self.do_classifier_free_guidance:
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+        if self.do_classifier_free_guidance and self.guidance_rescale > 0.0:
+            # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+            noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=self.guidance_rescale)
+
+        # compute the previous noisy sample x_t -> x_t-1
+        latents_dtype = latents.dtype
+        # latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+        # compute the previous noisy sample x_t -> x_t-1
+        latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=True)
+        try:
+            denoised = latents['pred_original_sample'] / self.vae.config.scaling_factor
+        except:
+            denoised = latents['denoised'] / self.vae.config.scaling_factor
+        latents = latents['prev_sample']
+
+        
+        # if latents.dtype != latents_dtype:
+        #     if torch.backends.mps.is_available():
+        #         # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+        latents = latents.to(self.vae.dtype)
+        denoised = denoised.to(self.vae.dtype)
+        
+        if i in save_timesteps:
+            if use_clip:
+                denoised = self.vae.decode(denoised.to(self.vae.dtype), return_dict=False)[0]
+                denoised = F.adaptive_avg_pool2d(denoised, (224, 224))
+                denoised = normalize_image(denoised)
+                if 'dino' in encoder:
+                    denoised = clip(denoised)
+                    denoised = denoised.pooler_output
+                    denoised = denoised.cpu().view(denoised.shape[0], -1)
+                else:
+                    denoised = clip.get_image_features(denoised)
+                    denoised = denoised.cpu().view(denoised.shape[0], -1)
+                    
+                # denoised = clip.get_image_features(denoised)
+            clip_features.append(denoised)
+
+        
+        
+
+        if callback_on_step_end is not None:
+            callback_kwargs = {}
+            for k in callback_on_step_end_tensor_inputs:
+                callback_kwargs[k] = locals()[k]
+            callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+            latents = callback_outputs.pop("latents", latents)
+            prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+            negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+            add_text_embeds = callback_outputs.pop("add_text_embeds", add_text_embeds)
+            negative_pooled_prompt_embeds = callback_outputs.pop(
+                "negative_pooled_prompt_embeds", negative_pooled_prompt_embeds
+            )
+            add_time_ids = callback_outputs.pop("add_time_ids", add_time_ids)
+            negative_add_time_ids = callback_outputs.pop("negative_add_time_ids", negative_add_time_ids)
+
+        # call the callback, if provided
+        if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+            # progress_bar.update()
+            if callback is not None and i % callback_steps == 0:
+                step_idx = i // getattr(self.scheduler, "order", 1)
+                callback(step_idx, t, latents)
+
+        if XLA_AVAILABLE:
+            xm.mark_step()
+
+    if not output_type == "latent":
+        # make sure the VAE is in float32 mode, as it overflows in float16
+        needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
+
+        if needs_upcasting:
+            self.upcast_vae()
+            latents = latents.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)
+        elif latents.dtype != self.vae.dtype:
+            if torch.backends.mps.is_available():
+                # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                self.vae = self.vae.to(latents.dtype)
+
+        # unscale/denormalize the latents
+        # denormalize with the mean and std if available and not None
+        has_latents_mean = hasattr(self.vae.config, "latents_mean") and self.vae.config.latents_mean is not None
+        has_latents_std = hasattr(self.vae.config, "latents_std") and self.vae.config.latents_std is not None
+        if has_latents_mean and has_latents_std:
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+            )
+            latents_std = (
+                torch.tensor(self.vae.config.latents_std).view(1, 4, 1, 1).to(latents.device, latents.dtype)
+            )
+            latents = latents * latents_std / self.vae.config.scaling_factor + latents_mean
+        else:
+            latents = latents / self.vae.config.scaling_factor
+
+        image = self.vae.decode(latents, return_dict=False)[0]
+
+        # cast back to fp16 if needed
+        if needs_upcasting:
+            self.vae.to(dtype=torch.float16)
+    else:
+        image = latents
+
+    if not output_type == "latent":
+
+        image = self.image_processor.postprocess(image, output_type=output_type)
+
+    # Offload all models
+    self.maybe_free_model_hooks()
+
+    return image, clip_features
+
+@torch.no_grad()
+
+def call_flux(
+    self,
+    prompt: Union[str, List[str]] = None,
+    prompt_2: Optional[Union[str, List[str]]] = None,
+    height: Optional[int] = None,
+    width: Optional[int] = None,
+    num_inference_steps: int = 28,
+    timesteps: List[int] = None,
+    guidance_scale: float = 7.0,
+    num_images_per_prompt: Optional[int] = 1,
+    generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    latents: Optional[torch.FloatTensor] = None,
+    prompt_embeds: Optional[torch.FloatTensor] = None,
+    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+    output_type: Optional[str] = "pil",
+    return_dict: bool = True,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+    callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+    max_sequence_length: int = 512,
+    verbose=False,
+    save_timesteps = None,
+    clip=None,
+    use_clip=True,
+    encoder='clip'
+):
+    
+
+    height = height or self.default_sample_size * self.vae_scale_factor
+    width = width or self.default_sample_size * self.vae_scale_factor
+
+    # 1. Check inputs. Raise error if not correct
+    self.check_inputs(
+        prompt,
+        prompt_2,
+        height,
+        width,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        max_sequence_length=max_sequence_length,
+    )
+
+    self._guidance_scale = guidance_scale
+    self._joint_attention_kwargs = joint_attention_kwargs
+    self._interrupt = False
+
+    # 2. Define call parameters
+    if prompt is not None and isinstance(prompt, str):
+        batch_size = 1
+    elif prompt is not None and isinstance(prompt, list):
+        batch_size = len(prompt)
+    else:
+        batch_size = prompt_embeds.shape[0]
+
+    device = self._execution_device
+
+    lora_scale = (
+        self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+    )
+    (
+        prompt_embeds,
+        pooled_prompt_embeds,
+        text_ids,
+    ) = self.encode_prompt(
+        prompt=prompt,
+        prompt_2=prompt_2,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_prompt_embeds,
+        device=device,
+        num_images_per_prompt=num_images_per_prompt,
+        max_sequence_length=max_sequence_length,
+        lora_scale=lora_scale,
+    )
+
+    # 4. Prepare latent variables
+    num_channels_latents = self.transformer.config.in_channels // 4
+    latents, latent_image_ids = self.prepare_latents(
+        batch_size * num_images_per_prompt,
+        num_channels_latents,
+        height,
+        width,
+        prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+
+    # 5. Prepare timesteps
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+    image_seq_len = latents.shape[1]
+    mu = calculate_shift(
+        image_seq_len,
+        self.scheduler.config.base_image_seq_len,
+        self.scheduler.config.max_image_seq_len,
+        self.scheduler.config.base_shift,
+        self.scheduler.config.max_shift,
+    )
+    timesteps, num_inference_steps = retrieve_timesteps(
+        self.scheduler,
+        num_inference_steps,
+        device,
+        timesteps,
+        sigmas,
+        mu=mu,
+    )
+
+    timesteps = timesteps
+    num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+    self._num_timesteps = len(timesteps)
+
+    # handle guidance
+    if self.transformer.config.guidance_embeds:
+        guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+        guidance = guidance.expand(latents.shape[0])
+    else:
+        guidance = None
+    clip_features = []
+    # 6. Denoising loop
+    with self.progress_bar(total=num_inference_steps) as progress_bar:
+        for i, t in enumerate(timesteps):
+            if self.interrupt:
+                continue
+
+            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+           
+            noise_pred = self.transformer(
+                hidden_states=latents,
+                # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
+                timestep=timestep / 1000,
+                guidance=guidance,
+                pooled_projections=pooled_prompt_embeds,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+                joint_attention_kwargs=self.joint_attention_kwargs,
+                return_dict=False,
+            )[0]
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents_dtype = latents.dtype
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(noise_pred, t, latents, return_dict=True)
+
+ 
+            denoised = latents['prev_sample'] 
+            latents = latents['prev_sample']
+
+            denoised = self._unpack_latents(denoised, height, width, self.vae_scale_factor)
+            denoised = (denoised / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            denoised = self.vae.decode(denoised, return_dict=False)[0]
+            denoised = F.adaptive_avg_pool2d(denoised, (224, 224))
+            if 'dino' in encoder:
+                outputs = clip(**inputs)
+                denoised = outputs.pooler_output
+                denoised = denoised.cpu().view(denoised.shape[0], -1)
+            else:
+                denoised = clip.get_image_features(denoised)
+                denoised = denoised.cpu().view(denoised.shape[0], -1)
+
+            clip_features.append()
+           
+            if latents.dtype != latents_dtype:
+                if torch.backends.mps.is_available():
+                    # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                    latents = latents.to(latents_dtype)
+
+            if callback_on_step_end is not None:
+                callback_kwargs = {}
+                for k in callback_on_step_end_tensor_inputs:
+                    callback_kwargs[k] = locals()[k]
+                callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                latents = callback_outputs.pop("latents", latents)
+                prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+            # call the callback, if provided
+            if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                progress_bar.update()
+
+            if XLA_AVAILABLE:
+                xm.mark_step()
+
+    if output_type == "latent":
+        image = latents
+        return image
+
+    else:
+        latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+        image = self.vae.decode(latents, return_dict=False)[0]
+        image = self.image_processor.postprocess(image, output_type=output_type)
+
+    # Offload all models
+    self.maybe_free_model_hooks()
+
+    if not return_dict:
+        return (image,)
+
+    return image, clip_features
+
+
+
+
+def get_diffusion_clip_directions(prompts, unet, tokenizers, text_encoders, vae, noise_scheduler, clip, batchsize=1, height=1024, width=1024, max_denoising_steps=4, savepath_training_images=None, use_clip=True,encoder='clip'):
+    device = unet.device
+    vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
+
+    os.makedirs(savepath_training_images, exist_ok=True)
+
+
+    if len(noise_scheduler.timesteps) != max_denoising_steps:
+        noise_scheduler_orig = noise_scheduler
+        max_denoising_steps_orig = len(noise_scheduler.timesteps)
+        noise_scheduler.set_timesteps(max_denoising_steps)
+        timesteps_distilled = noise_scheduler.timesteps
+        
+        noise_scheduler.set_timesteps(max_denoising_steps_orig)
+        timesteps_full = noise_scheduler.timesteps
+        save_timesteps = []
+        for timesteps_to_distilled in range(max_denoising_steps):
+            # Get the value from timesteps_distilled that we want to find in timesteps_full
+            value_to_find = timesteps_distilled[timesteps_to_distilled]
+            timesteps_to_full = (timesteps_full == value_to_find).nonzero().item()
+            save_timesteps.append(timesteps_to_full)
+
+        guidance_scale = 7
+    else:
+        max_denoising_steps_orig = max_denoising_steps
+        save_timesteps = [i for i in range(max_denoising_steps_orig)]
+        guidance_scale = 7
+        if max_denoising_steps_orig <=4:
+            guidance_scale = 0
+        
+    noise_scheduler.set_timesteps(max_denoising_steps_orig)
+    # if max_denoising_steps_orig == 1:
+    #     noise_scheduler.set_timesteps(timesteps=[399],
+    #                                  device=device)
+    
+    weight_dtype = unet.dtype
+    device = unet.device 
+    StableDiffusionXLPipeline.__call__ = call_sdxl
+    pipe = StableDiffusionXLPipeline(vae = vae,
+        text_encoder= text_encoders[0],
+        text_encoder_2=text_encoders[1],
+        tokenizer = tokenizers[0],
+        tokenizer_2= tokenizers[1],
+        unet=unet,
+        scheduler=noise_scheduler)
+    pipe.to(unet.device)
+    # print(guidance_scale, max_denoising_steps_orig, save_timesteps)
+    images, clip_features = pipe(prompts, guidance_scale=guidance_scale, num_inference_steps = max_denoising_steps_orig, clip=clip, save_timesteps =save_timesteps, use_clip=use_clip, encoder=encoder)
+    
+    return images, torch.stack(clip_features)
+
+
+
+def get_flux_clip_directions(prompts, transformer, tokenizers, text_encoders, vae, noise_scheduler, clip, batchsize=1, height=1024, width=1024, max_denoising_steps=4, savepath_training_images=None, use_clip=True):
+    device = transformer.device
+    FluxPipeline.__call__ = call_flux
+    pipe = FluxPipeline(noise_scheduler,
+                    vae,
+                    text_encoders[0],
+                    tokenizers[0],
+                    text_encoders[1],
+                    tokenizers[1],
+                    transformer,
+                   )
+    pipe.set_progress_bar_config(disable=True)
+
+    os.makedirs(savepath_training_images, exist_ok=True)
+
+    images, clip_features = pipe(
+        prompts,
+        height=height,
+        width=width,
+        guidance_scale=0,
+        num_inference_steps=4,
+        max_sequence_length=256,
+        num_images_per_prompt=1,
+        output_type='pil',
+        clip=clip
+    )
+    
+    return images, torch.stack(clip_features)
+
+
+
+
+def get_diffusion_clip_directions(prompts, unet, tokenizers, text_encoders, vae, noise_scheduler, clip, batchsize=1, height=1024, width=1024, max_denoising_steps=4, savepath_training_images=None, use_clip=True,encoder='clip', num_images_per_prompt=1):
+
+    
+    device = unet.device
+    vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
+    image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor)
+    os.makedirs(savepath_training_images, exist_ok=True)
+
+
+    if len(noise_scheduler.timesteps) != max_denoising_steps:
+        noise_scheduler_orig = noise_scheduler
+        max_denoising_steps_orig = len(noise_scheduler.timesteps)
+        noise_scheduler.set_timesteps(max_denoising_steps)
+        timesteps_distilled = noise_scheduler.timesteps
+        
+        noise_scheduler.set_timesteps(max_denoising_steps_orig)
+        timesteps_full = noise_scheduler.timesteps
+        save_timesteps = []
+        for timesteps_to_distilled in range(max_denoising_steps):
+            # Get the value from timesteps_distilled that we want to find in timesteps_full
+            value_to_find = timesteps_distilled[timesteps_to_distilled]
+            timesteps_to_full = (timesteps_full == value_to_find).nonzero().item()
+            save_timesteps.append(timesteps_to_full)
+
+        guidance_scale = 7
+    else:
+        max_denoising_steps_orig = max_denoising_steps
+        save_timesteps = [i for i in range(max_denoising_steps_orig)]
+        guidance_scale = 7
+        if max_denoising_steps_orig <=4:
+            guidance_scale = 0
+        
+    noise_scheduler.set_timesteps(max_denoising_steps_orig)
+    # if max_denoising_steps_orig == 1:
+    #     noise_scheduler.set_timesteps(timesteps=[399],
+    #                                  device=device)
+    
+    weight_dtype = unet.dtype
+    device = unet.device 
+    StableDiffusionXLPipeline.__call__ = call_sdxl
+    pipe = StableDiffusionXLPipeline(vae = vae,
+        text_encoder= text_encoders[0],
+        text_encoder_2=text_encoders[1],
+        tokenizer = tokenizers[0],
+        tokenizer_2= tokenizers[1],
+        unet=unet,
+        scheduler=noise_scheduler)
+    pipe.to(unet.device)
+    # print(guidance_scale, max_denoising_steps_orig, save_timesteps)
+    images, clip_features = pipe(prompts, guidance_scale=guidance_scale, num_inference_steps = max_denoising_steps_orig, clip=clip, save_timesteps =save_timesteps, use_clip=use_clip, encoder=encoder)
+    
+    return images, torch.stack(clip_features)
+
+
+
+