utils/ai_generator_diffusers_flux.py

# utils/ai_generator_diffusers_flux.py
import os
import utils.constants as constants
import spaces
import torch
from diffusers import FluxPipeline,FluxImg2ImgPipeline,FluxControlPipeline
import accelerate 
import transformers
import safetensors
import xformers
from diffusers.utils import load_image
from huggingface_hub import hf_hub_download
from PIL import Image
from tempfile import NamedTemporaryFile
from src.condition import Condition
from utils.image_utils import (
     crop_and_resize_image,
)
from utils.version_info import (
    get_torch_info,
    get_diffusers_version,
    get_transformers_version,
    get_xformers_version
)
from utils.lora_details import get_trigger_words, approximate_token_count, split_prompt_precisely
from utils.color_utils import detect_color_format
import utils.misc as misc
from pathlib import Path
import warnings

warnings.filterwarnings("ignore", message=".*Torch was not compiled with flash attention.*")
#print(torch.__version__)  # Ensure it's 2.0 or newer
#print(torch.cuda.is_available())  # Ensure CUDA is available

PIPELINE_CLASSES = {
    "FluxPipeline": FluxPipeline,
    "FluxImg2ImgPipeline": FluxImg2ImgPipeline
}

@spaces.GPU(duration=140)
def generate_image_from_text(
    text,
    model_name="black-forest-labs/FLUX.1-dev",
    lora_weights=None,
    conditioned_image=None,
    image_width=1344,
    image_height=848,
    guidance_scale=3.5,
    num_inference_steps=50,
    seed=0,
    additional_parameters=None
):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"device:{device}\nmodel_name:{model_name}\n")
    pipe = FluxPipeline.from_pretrained(
        model_name,
        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32
    ).to(device)
    pipe = pipe.to(device)
    pipe.enable_model_cpu_offload()
    # Load and apply LoRA weights
    if lora_weights:
        for lora_weight in lora_weights:
            lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
            if lora_configs:
                for config in lora_configs:
                    weight_name = config.get("weight_name")
                    adapter_name = config.get("adapter_name")
                    pipe.load_lora_weights(
                        lora_weight,
                        weight_name=weight_name,
                        adapter_name=adapter_name,
                        use_auth_token=constants.HF_API_TOKEN
                    )
            else:
                pipe.load_lora_weights(lora_weight, use_auth_token=constants.HF_API_TOKEN)
    generator = torch.Generator(device=device).manual_seed(seed)
    conditions = []
    if conditioned_image is not None:
        conditioned_image = crop_and_resize_image(conditioned_image, 1024, 1024)
        condition = Condition("subject", conditioned_image)
        conditions.append(condition)
    generate_params = {
        "prompt": text,
        "height": image_height,
        "width": image_width,
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps,
        "generator": generator,
        "conditions": conditions if conditions else None
    }
    if additional_parameters:
        generate_params.update(additional_parameters)
    generate_params = {k: v for k, v in generate_params.items() if v is not None}
    result = pipe(**generate_params)
    image = result.images[0]
    pipe.unload_lora_weights()
    return image

@spaces.GPU(duration=140)
def generate_image_lowmem(
    text,
    neg_prompt=None,
    model_name="black-forest-labs/FLUX.1-dev",
    lora_weights=None,
    conditioned_image=None,
    image_width=1368,
    image_height=848,
    guidance_scale=3.5,
    num_inference_steps=30,
    seed=0,
    true_cfg_scale=1.0,
    pipeline_name="FluxPipeline",
    strength=0.75,
    additional_parameters=None
):
    print(f"\n {get_torch_info()}\n")
    # Retrieve the pipeline class from the mapping
    pipeline_class = PIPELINE_CLASSES.get(pipeline_name)
    if not pipeline_class:
        raise ValueError(f"Unsupported pipeline type '{pipeline_name}'. "
                        f"Available options: {list(PIPELINE_CLASSES.keys())}")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"device:{device}\nmodel_name:{model_name}\nlora_weights:{lora_weights}\n")
    
    # Disable gradient calculations
    with torch.no_grad():
        # Initialize the pipeline inside the context manager
        pipe = pipeline_class.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16 if device == "cuda" else torch.bfloat32            
        ).to(device)
        # Optionally, don't use CPU offload if not necessary
        pipe.enable_model_cpu_offload()
        # alternative version that may be more efficient
        # pipe.enable_sequential_cpu_offload()
        flash_attention_enabled = torch.backends.cuda.flash_sdp_enabled()
        if flash_attention_enabled == False:
            #Enable xFormers memory-efficient attention (optional)
            pipe.enable_xformers_memory_efficient_attention()
            print("\nEnabled xFormers memory-efficient attention.\n")
        else:            
            pipe.attn_implementation="flash_attention_2"
            print("\nEnabled flash_attention_2.\n")
        if pipeline_name == "FluxPipeline":
            pipe.enable_vae_tiling()
        # Load LoRA weights
        # note: does not yet handle multiple LoRA weights with different names, needs .set_adapters(["depth", "hyper-sd"], adapter_weights=[0.85, 0.125])
        if lora_weights:
            for lora_weight in lora_weights:
                lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
                lora_weight_set = False
                if lora_configs:
                    for config in lora_configs:
                        # Load LoRA weights with optional weight_name and adapter_name
                        if 'weight_name' in config:
                            weight_name = config.get("weight_name")
                            adapter_name = config.get("adapter_name")
                            lora_collection = config.get("lora_collection")
                            if weight_name and adapter_name and lora_collection and lora_weight_set == False:
                                pipe.load_lora_weights(
                                    lora_collection,
                                    weight_name=weight_name,
                                    adapter_name=adapter_name,
                                    token=constants.HF_API_TOKEN
                                )
                                lora_weight_set = True
                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
                            elif weight_name and adapter_name==None and lora_collection and lora_weight_set == False:
                                pipe.load_lora_weights(
                                    lora_collection,
                                    weight_name=weight_name,
                                    token=constants.HF_API_TOKEN
                                )
                                lora_weight_set = True
                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}, lora_collection={lora_collection}\n")  
                            elif weight_name and adapter_name and lora_weight_set == False:
                                pipe.load_lora_weights(
                                    lora_weight,
                                    weight_name=weight_name,
                                    adapter_name=adapter_name,
                                    token=constants.HF_API_TOKEN
                                )
                                lora_weight_set = True
                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
                            elif weight_name and adapter_name==None and lora_weight_set == False:
                                pipe.load_lora_weights(
                                    lora_weight,
                                    weight_name=weight_name,
                                    token=constants.HF_API_TOKEN
                                )
                                lora_weight_set = True
                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
                            elif lora_weight_set == False:
                                pipe.load_lora_weights(
                                    lora_weight,
                                    token=constants.HF_API_TOKEN
                                )  
                                lora_weight_set = True
                                print(f"\npipe.load_lora_weights({lora_weight}, weight_name={weight_name}, adapter_name={adapter_name}\n")  
                        # Apply 'pipe' configurations if present
                        if 'pipe' in config:
                            pipe_config = config['pipe']
                            for method_name, params in pipe_config.items():
                                method = getattr(pipe, method_name, None)
                                if method:
                                    print(f"Applying pipe method: {method_name} with params: {params}")
                                    method(**params)
                                else:
                                    print(f"Method {method_name} not found in pipe.")
                else:
                    pipe.load_lora_weights(lora_weight, use_auth_token=constants.HF_API_TOKEN)
        # Set the random seed for reproducibility
        generator = torch.Generator(device=device).manual_seed(seed)
        conditions = []
        if conditioned_image is not None:
            conditioned_image = crop_and_resize_image(conditioned_image, image_width, image_height)
            condition = Condition("subject", conditioned_image)
            conditions.append(condition)
            print(f"\nAdded conditioned image.\n {conditioned_image.size}")
            # Prepare the parameters for image generation
            additional_parameters ={
                "strength": strength,
                "image": conditioned_image,
            }
        else:
            print("\nNo conditioned image provided.")
            if neg_prompt!=None:
                true_cfg_scale=1.1
            additional_parameters ={
                "negative_prompt": neg_prompt,
                "true_cfg_scale": true_cfg_scale,
            }
        # handle long prompts by splitting them
        if approximate_token_count(text) > 76:
            prompt, prompt2 = split_prompt_precisely(text)
            prompt_parameters = {
                "prompt" : prompt,
                "prompt_2": prompt2
            }
        else:
            prompt_parameters = {
                "prompt" :text
        }
        additional_parameters.update(prompt_parameters)
        # Combine all parameters
        generate_params = {
            "height": image_height,
            "width": image_width,
            "guidance_scale": guidance_scale,
            "num_inference_steps": num_inference_steps,
            "generator": generator,        }
        if additional_parameters:
            generate_params.update(additional_parameters)
        generate_params = {k: v for k, v in generate_params.items() if v is not None}
        print(f"generate_params: {generate_params}")
        # Generate the image
        result = pipe(**generate_params)
        image = result.images[0]
        # Clean up
        del result
        del conditions
        del generator
    # Delete the pipeline and clear cache
    del pipe
    torch.cuda.empty_cache()
    torch.cuda.ipc_collect()
    print(torch.cuda.memory_summary(device=None, abbreviated=False))
    return image

def generate_ai_image_local (
    map_option,
    prompt_textbox_value,
    neg_prompt_textbox_value,
    model="black-forest-labs/FLUX.1-dev",
    lora_weights=None,
    conditioned_image=None,
    height=512,
    width=912,
    num_inference_steps=30,
    guidance_scale=3.5,
    seed=777,
    pipeline_name="FluxPipeline",
    strength=0.75,
):
    try:
        if map_option != "Prompt":
            prompt = constants.PROMPTS[map_option]
            negative_prompt = constants.NEGATIVE_PROMPTS.get(map_option, "")
        else:
            prompt = prompt_textbox_value
            negative_prompt = neg_prompt_textbox_value or ""
        #full_prompt = f"{prompt} {negative_prompt}"
        additional_parameters = {}
        if lora_weights:
            for lora_weight in lora_weights:
                lora_configs = constants.LORA_DETAILS.get(lora_weight, [])
                for config in lora_configs:
                    if 'parameters' in config:
                        additional_parameters.update(config['parameters'])
                    elif 'trigger_words' in config:
                        trigger_words = get_trigger_words(lora_weight)
                        prompt = f"{trigger_words} {prompt}"
        for key, value in additional_parameters.items():
            if key in ['height', 'width', 'num_inference_steps', 'max_sequence_length']:
                additional_parameters[key] = int(value)
            elif key in ['guidance_scale','true_cfg_scale']:
                additional_parameters[key] = float(value)
        height = additional_parameters.get('height', height)
        width = additional_parameters.get('width', width)
        num_inference_steps = additional_parameters.get('num_inference_steps', num_inference_steps)
        guidance_scale = additional_parameters.get('guidance_scale', guidance_scale)
        print("Generating image with the following parameters:")
        print(f"Model: {model}")
        print(f"LoRA Weights: {lora_weights}")
        print(f"Prompt: {prompt}")
        print(f"Neg Prompt: {negative_prompt}")
        print(f"Height: {height}")
        print(f"Width: {width}")
        print(f"Number of Inference Steps: {num_inference_steps}")
        print(f"Guidance Scale: {guidance_scale}")
        print(f"Seed: {seed}")
        print(f"Additional Parameters: {additional_parameters}")
        print(f"Conditioned Image: {conditioned_image}")
        print(f"pipeline: {pipeline_name}")
        image = generate_image_lowmem(
            text=prompt,
            model_name=model,
            neg_prompt=negative_prompt,
            lora_weights=lora_weights,
            conditioned_image=conditioned_image,
            image_width=width,
            image_height=height,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            seed=seed,
            pipeline_name=pipeline_name,
            additional_parameters=additional_parameters
        )
        with NamedTemporaryFile(delete=False, suffix=".png") as tmp:
            image.save(tmp.name, format="PNG")
            constants.temp_files.append(tmp.name)
            print(f"Image saved to {tmp.name}")
            return tmp.name
    except Exception as e:
        print(f"Error generating AI image: {e}")
        return None

# does not work
#@spaces.GPU(duration=256)
def merge_LoRA_weights(model="black-forest-labs/FLUX.1-dev",
    lora_weights="Borcherding/FLUX.1-dev-LoRA-FractalLand-v0.1"):
    
    model_suffix = model.split("/")[-1]
    if model_suffix not in lora_weights:
        raise ValueError(f"The model suffix '{model_suffix}' must be in the lora_weights string '{lora_weights}' to proceed.")
    
    pipe = FluxPipeline.from_pretrained(model, torch_dtype=torch.bfloat16)
    pipe.load_lora_weights(lora_weights)
    pipe.save_lora_weights(os.getenv("TMPDIR"))
    lora_name = lora_weights.split("/")[-1] + "-merged"
    pipe.save_pretrained(lora_name)
    pipe.unload_lora_weights()