gradio_app.py

import spaces
import gradio as gr
import torch
import numpy as np
from PIL import Image
from accelerate import Accelerator
import os
import time
from torchvision import transforms
from safetensors.torch import load_file
from networks import lora_flux
from library import flux_utils, flux_train_utils_recraft as flux_train_utils, strategy_flux
import logging
from huggingface_hub import login
from huggingface_hub import hf_hub_download


# Set up logger
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG)

accelerator = Accelerator(mixed_precision='bf16', device_placement=True)

hf_token = os.getenv("HF_TOKEN")
login(token=hf_token)

# Model paths dynamically retrieved using selected model
model_paths = {
    'Wood Sculpture': {
        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
        'BASE_FILE': "flux1-dev-fp8.safetensors",
        'LORA_REPO': "showlab/makeanything",
        'LORA_FILE': "recraft/recraft_4f_wood_sculpture.safetensors"
    },
    'LEGO': {
        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
        'BASE_FILE': "flux1-dev-fp8.safetensors",
        'LORA_REPO': "showlab/makeanything",
        'LORA_FILE': "recraft/recraft_9f_lego.safetensors"
    },
    'Sketch': {
        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
        'BASE_FILE': "flux1-dev-fp8.safetensors",
        'LORA_REPO': "showlab/makeanything",
        'LORA_FILE': "recraft/recraft_9f_sketch.safetensors"
    },
    'Portrait': {
        'BASE_FLUX_CHECKPOINT': "Kijai/flux-fp8",
        'BASE_FILE': "flux1-dev-fp8.safetensors",
        'LORA_REPO': "showlab/makeanything",
        'LORA_FILE': "recraft/recraft_9f_portrait.safetensors"
    }
}

# Common paths
clip_repo_id = "comfyanonymous/flux_text_encoders"
t5xxl_file = "t5xxl_fp8_e4m3fn.safetensors"
clip_l_file = "clip_l.safetensors"
ae_repo_id = "black-forest-labs/FLUX.1-dev"
ae_file = "ae.safetensors"


# Model placeholders
model = None
clip_l = None
t5xxl = None
ae = None
lora_model = None

# Function to load a file from Hugging Face Hub
def download_file(repo_id, file_name):
    return hf_hub_download(repo_id=repo_id, filename=file_name)

# Load model function
def load_target_model(selected_model):
    global model, clip_l, t5xxl, ae, lora_model

    # Fetch paths based on the selected model
    model_path = model_paths[selected_model]
    base_checkpoint_repo = model_path['BASE_FLUX_CHECKPOINT']
    base_checkpoint_file = model_path['BASE_FILE']
    lora_repo = model_path['LORA_REPO']
    lora_file = model_path['LORA_FILE']

    # Download necessary files
    BASE_FLUX_CHECKPOINT = download_file(base_checkpoint_repo, base_checkpoint_file)
    CLIP_L_PATH = download_file(clip_repo_id, clip_l_file)
    T5XXL_PATH = download_file(clip_repo_id, t5xxl_file)
    AE_PATH = download_file(ae_repo_id, ae_file)
    LORA_WEIGHTS_PATH = download_file(lora_repo, lora_file)

    logger.info("Loading models...")
    try:
        _, model = flux_utils.load_flow_model(
            BASE_FLUX_CHECKPOINT, torch.float8_e4m3fn, "cpu", disable_mmap=False
        )
        clip_l = flux_utils.load_clip_l(CLIP_L_PATH, torch.bfloat16, "cpu", disable_mmap=False)
        clip_l.eval()
        t5xxl = flux_utils.load_t5xxl(T5XXL_PATH, torch.bfloat16, "cpu", disable_mmap=False)
        t5xxl.eval()
        ae = flux_utils.load_ae(AE_PATH, torch.bfloat16, "cpu", disable_mmap=False)
        logger.info("Models loaded successfully.")
        return model, [clip_l, t5xxl], ae
    except Exception as e:
        logger.error(f"Error loading models: {e}")
        raise

# Image pre-processing (resize and padding)
class ResizeWithPadding:
    def __init__(self, size, fill=255):
        self.size = size
        self.fill = fill

    def __call__(self, img):
        if isinstance(img, np.ndarray):
            img = Image.fromarray(img)
        elif not isinstance(img, Image.Image):
            raise TypeError("Input must be a PIL Image or a NumPy array")

        width, height = img.size

        # Convert to RGB to remove transparency, fill with white background if necessary
        if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
            background = Image.new("RGB", img.size, (fill, fill, fill))
            background.paste(img, mask=img.split()[-1])  # Use alpha channel as mask
            img = background

        if width == height:
            img = img.resize((self.size, self.size), Image.LANCZOS)
        else:
            max_dim = max(width, height)
            new_img = Image.new("RGB", (max_dim, max_dim), (self.fill, self.fill, self.fill))
            new_img.paste(img, ((max_dim - width) // 2, (max_dim - height) // 2))
            img = new_img.resize((self.size, self.size), Image.LANCZOS)
        return img

# The function to generate image from a prompt and conditional image
@spaces.GPU(duration=180)
def infer(prompt, sample_image, frame_num, seed=0):
    global model, clip_l, t5xxl, ae, lora_model
    if model is None or lora_model is None or clip_l is None or t5xxl is None or ae is None:
        logger.error("Models not loaded. Please load the models first.")
        return None

    logger.info(f"Started generating image with prompt: {prompt}")

    lora_model.to("cuda")
    
    model.eval()
    clip_l.eval()
    t5xxl.eval()
    ae.eval()

    # # Load models
    # model, [clip_l, t5xxl], ae = load_target_model()

    # # LoRA
    # multiplier = 1.0
    # weights_sd = load_file(LORA_WEIGHTS_PATH)
    # lora_model, _ = lora_flux.create_network_from_weights(multiplier, None, ae, [clip_l, t5xxl], model, weights_sd,
    #                                                       True)

    # lora_model.apply_to([clip_l, t5xxl], model)
    # info = lora_model.load_state_dict(weights_sd, strict=True)
    # logger.info(f"Loaded LoRA weights from {LORA_WEIGHTS_PATH}: {info}")
    # lora_model.eval()
    # lora_model.to(device)

    logger.debug(f"Using seed: {seed}")

    # Preprocess the conditional image
    resize_transform = ResizeWithPadding(size=512) if frame_num == 4 else ResizeWithPadding(size=352)
    img_transforms = transforms.Compose([
        resize_transform,
        transforms.ToTensor(),
        transforms.Normalize([0.5], [0.5]),
    ])
    image = img_transforms(np.array(sample_image, dtype=np.uint8)).unsqueeze(0).to(
        device=device,
        dtype=torch.bfloat16
    )
    logger.debug("Conditional image preprocessed.")

    # Encode the image to latents
    ae.to(device)
    latents = ae.encode(image)
    logger.debug("Image encoded to latents.")

    conditions = {}
    conditions[prompt] = latents.to("cpu")

    ae.to("cpu")
    clip_l.to(device)
    t5xxl.to(device)

    # Encode the prompt
    tokenize_strategy = strategy_flux.FluxTokenizeStrategy(512)
    text_encoding_strategy = strategy_flux.FluxTextEncodingStrategy(True)
    tokens_and_masks = tokenize_strategy.tokenize(prompt)
    l_pooled, t5_out, txt_ids, t5_attn_mask = text_encoding_strategy.encode_tokens(tokenize_strategy, [clip_l, t5xxl], tokens_and_masks, True)

    logger.debug("Prompt encoded.")

    # Prepare the noise and other parameters
    width = 1024 if frame_num == 4 else 1056
    height = 1024 if frame_num == 4 else 1056

    height = max(64, height - height % 16)
    width = max(64, width - width % 16)

    packed_latent_height = height // 16
    packed_latent_width = width // 16
    
    noise = torch.randn(1, packed_latent_height * packed_latent_width, 16 * 2 * 2, device=device, dtype=torch.float16)
    logger.debug("Noise prepared.")

    # Generate the image
    timesteps = flux_train_utils.get_schedule(20, noise.shape[1], shift=True)  # Sample steps = 20
    img_ids = flux_utils.prepare_img_ids(1, packed_latent_height, packed_latent_width).to(device)

    t5_attn_mask = t5_attn_mask.to(device)
    ae_outputs = conditions[prompt]

    logger.debug("Image generation parameters set.")
    
    args = lambda: None
    args.frame_num = frame_num

    clip_l.to("cpu")
    t5xxl.to("cpu")

    model.to(device)

    # Run the denoising process
    with accelerator.autocast(), torch.no_grad():
        x = flux_train_utils.denoise(
            args, model, noise, img_ids, t5_out, txt_ids, l_pooled, timesteps=timesteps, guidance=1.0, t5_attn_mask=t5_attn_mask, ae_outputs=ae_outputs
        )
    logger.debug("Denoising process completed.")

    # Decode the final image
    x = x.float()
    x = flux_utils.unpack_latents(x, packed_latent_height, packed_latent_width)
    model.to("cpu")
    ae.to(device)
    with accelerator.autocast(), torch.no_grad():
        x = ae.decode(x)
    logger.debug("Latents decoded into image.")
    ae.to("cpu")

    # Convert the tensor to an image
    x = x.clamp(-1, 1)
    x = x.permute(0, 2, 3, 1)
    generated_image = Image.fromarray((127.5 * (x + 1.0)).float().cpu().numpy().astype(np.uint8)[0])

    logger.info("Image generation completed.")
    return generated_image

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## FLUX Image Generation")
    
    with gr.Row():
        # Dropdown for selecting the recraft model
        recraft_model = gr.Dropdown(
            label="Select Recraft Model", 
            choices=["Wood Sculpture", "LEGO", "Sketch", "Portrait"], 
            value="Wood Sculpture"
        )

        # Input for the prompt
        prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here", lines=1)

        # File upload for image
        sample_image = gr.Image(label="Upload a Conditional Image", type="pil")

        # Frame number selection
        frame_num = gr.Radio([4, 9], label="Select Frame Number", value=4)

        # Seed
        seed = gr.Slider(0, np.iinfo(np.int32).max, step=1, label="Seed", value=0)

        # Load Model Button
        load_button = gr.Button("Load Model")

        # Run Button
        run_button = gr.Button("Generate Image")

        # Output result
        result_image = gr.Image(label="Generated Image")

    # Load model button action
    load_button.click(fn=load_target_model, inputs=[recraft_model], outputs=[])

    # Run Button 
    run_button.click(fn=infer, inputs=[prompt, sample_image, frame_num, seed], outputs=[result_image])

    # Launch the Gradio app
    demo.launch()