Generate Image with LoRA

[alfredplpl]

I have create a LoRA file.
But, I cannot load the file by the following code.

pipeline.dit_model.load_lora_weights("/media/user/NVM/F-Lite-DiT/final_model")

I got the error:

Traceback (most recent call last):
  File "/mnt/raid0/github/f-lite/inference.py", line 14, in <module>
    pipeline.dit_model.load_lora_weights("/media/user/NVM/F-Lite-DiT/final_model")
  File "/mnt/raid0/github/f-lite/f_lite/model.py", line 379, in load_lora_weights
    set_peft_model_state_dict(self, lora_state_dict)
  File "/mnt/raid0/github/f-lite/.venv/lib/python3.12/site-packages/peft/utils/save_and_load.py", line 343, in set_peft_model_state_dict
    config = model.peft_config[adapter_name]
             ^^^^^^^^^^^^^^^^^
  File "/mnt/raid0/github/f-lite/.venv/lib/python3.12/site-packages/diffusers/models/modeling_utils.py", line 291, in __getattr__
    return super().__getattr__(name)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/mnt/raid0/github/f-lite/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1940, in __getattr__
    raise AttributeError(
AttributeError: 'DiT' object has no attribute 'peft_config'. Did you mean: 'from_config'?

How can I do?

[erwold]

Can you try this?

dit_model.add_adapter(lora_config)
lora_state_dict = {k: v for k, v in state_dict.items() if '.lora_A.' in k or '.lora_B.' in k}
set_peft_model_state_dict(dit_model, lora_state_dict)

[alfredplpl]

Thank you for reply.

Unfortunately, I cannot understand the your suggestion.

Please tell me by filling the following code:

import torch
from f_lite import FLitePipeline
from peft      import PeftModel

# Trick required because it is not a native diffusers model
from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMPORTABLE_CLASSES
LOADABLE_CLASSES["f_lite"] = LOADABLE_CLASSES["f_lite.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
ALL_IMPORTABLE_CLASSES["DiT"] = ["save_pretrained", "from_pretrained"]

pipeline = FLitePipeline.from_pretrained("/mnt/nvme2n1p2/F-Lite", torch_dtype=torch.bfloat16)
pipeline.enable_model_cpu_offload() # For less memory consumption. Alternatively, pipeline.to("cuda")
pipeline.vae.enable_slicing()
pipeline.vae.enable_tiling()

### fill here

# Generate an image
output = pipeline(
    prompt="test"
    height=1024,
    width=1024,
    num_inference_steps=30,
    guidance_scale=3.0,
    negative_prompt=None,
)

# Save the generated image
output.images[0].save("/mnt/nvme2n1p2/generated_image.png")

[erwold]

Here~

import torch
from f_lite import FLitePipeline
from peft import LoraConfig, set_peft_model_state_dict
from diffusers.pipelines.pipeline_loading_utils import LOADABLE_CLASSES, ALL_IMPORTABLE_CLASSES

# Register the DiT model with diffusers
LOADABLE_CLASSES["f_lite"] = LOADABLE_CLASSES["f_lite.model"] = {"DiT": ["save_pretrained", "from_pretrained"]}
ALL_IMPORTABLE_CLASSES["DiT"] = ["save_pretrained", "from_pretrained"]

# Load the base pipeline model
pipeline = FLitePipeline.from_pretrained(
    "/path/to/your/base_model", 
    torch_dtype=torch.bfloat16
)

# Enable memory optimizations
pipeline.enable_model_cpu_offload()  # For less memory consumption
pipeline.vae.enable_slicing()
pipeline.vae.enable_tiling()

# Set LoRA configuration
lora_weights_path = "/path/to/your/lora_checkpoint"
target_modules = ["qkv", "context_kv", "proj", "q"]  # Target modules detected from LoRA weights
lora_rank = 128  # LoRA rank of your lora weights

# Create LoRA configuration
lora_config = LoraConfig(
    r=lora_rank,
    lora_alpha=lora_rank,  # Typically set to the same value as rank
    target_modules=target_modules,
    bias="none",
    init_lora_weights=False,  # Don't initialize weights, use loaded ones
)

# Add LoRA adapter to the model
pipeline.dit_model.add_adapter(lora_config)

# Load LoRA weights
lora_state_dict = torch.load(f"{lora_weights_path}/lora_weights.pt", weights_only=True)
set_peft_model_state_dict(pipeline.dit_model, lora_state_dict)

# Generate an image
output = pipeline(
    prompt="a beautiful woman",  # Your text prompt
    height=1024,  # Output image height
    width=1024,   # Output image width
    num_inference_steps=30,  # Number of denoising steps
    guidance_scale=3.0,  # CFG scale - higher values adhere more strongly to prompt
    negative_prompt=None,  # Optional negative prompt
)

# Save the generated image
output.images[0].save("output.png")

if you don't know your lora config, use this script:

import torch

# Path to your LoRA weights
lora_weights_path = "/path/to/your/lora_checkpoint"
lora_file = f"{lora_weights_path}/lora_weights.pt"

# Load LoRA weights
print(f"Loading LoRA weights: {lora_file}")
lora_weights = torch.load(lora_file, weights_only=True)

# Print keys to identify target modules
keys = list(lora_weights.keys())
print(f"\nTotal keys in weights file: {len(keys)}")
print(f"\nFirst 5 keys as examples: {keys[:5]}")

# Extract target modules
target_modules = set()
for key in keys:
    if '.lora_A.' in key:
        parts = key.split('.')
        # Target module name is typically before lora_A
        for i, part in enumerate(parts):
            if part == 'lora_A':
                if i > 0:
                    target_modules.add(parts[i-1])

print(f"\nDetected target modules: {target_modules}")

# Determine LoRA rank
for key in keys:
    if '.lora_A.' in key and key.endswith('.weight'):
        shape = lora_weights[key].shape
        if len(shape) >= 2:
            # First dimension of A matrix is typically the rank
            lora_rank = shape[0]
            print(f"\nDetected LoRA rank: {lora_rank}")
            print(f"Matrix shape: {shape}")
            break

[alfredplpl]

Thanks.

I have applied the LoRA.

Before:

After:

I wanna make an adapter for UMT5-XXL because I wanna use the F-Lite in Japaenese.
The LoRA makes me happy to achieve the piece of the dream.

[alfredplpl]

And Japanese LoRA.

prompt: "a photo of Japanese sushi including tuna."

Before:

After:

LoRA: https://huggingface.co/alfredplpl/F-Lite-Japanese-LoRA