gradio_app_asy.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
import math
import json
from torchvision import transforms
from safetensors.torch import load_file
from networks import asylora_flux as lora_flux
from library import flux_utils, strategy_flux
import flux_minimal_inference_asylora as flux_train_utils
import logging
from huggingface_hub import login
from huggingface_hub import hf_hub_download

device = "cuda" if torch.cuda.is_available() else "cpu"

# 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': "showlab/makeanything",
#         'BASE_FILE': "flux_merge_lora/flux_merge_4f_wood-fp16.safetensors",
#         'LORA_REPO': "showlab/makeanything",
#         'LORA_FILE': "recraft/recraft_4f_wood_sculpture.safetensors",
#         "Frame": 4
#     },
#     'LEGO': {
#         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
#         'BASE_FILE': "flux_merge_lora/flux_merge_9f_lego-fp16.safetensors",
#         'LORA_REPO': "showlab/makeanything",
#         'LORA_FILE': "recraft/recraft_9f_lego.safetensors",
#         "Frame": 9
#     },
#     'Sketch': {
#         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
#         'BASE_FILE': "flux_merge_lora/flux_merge_9f_portrait-fp16.safetensors",
#         'LORA_REPO': "showlab/makeanything",
#         'LORA_FILE': "recraft/recraft_9f_sketch.safetensors",
#         "Frame": 9
#     },
#     'Portrait': {
#         'BASE_FLUX_CHECKPOINT': "showlab/makeanything",
#         'BASE_FILE': "flux_merge_lora/flux_merge_9f_sketch-fp16.safetensors",
#         'LORA_REPO': "showlab/makeanything",
#         'LORA_FILE': "recraft/recraft_9f_portrait.safetensors",
#         "Frame": 9
#     }
# }

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

domain_index = {
    'LEGO': 1, 'Cook': 2, 'Painting': 3, 'Icon': 4, 'Landscape illustration': 5, 
    'Portrait': 6, 'Transformer': 7, 'Sand art': 8, 'Illustration': 9,  'Sketch': 10,
    'Clay toys': 11, 'Clay sculpture': 12, 'Zbrush Modeling': 13, 'Wood sculpture': 14,
    'Ink painting': 15, 'Pencil sketch': 16, 'Fabric toys': 17, 'Oil painting': 18,
    'Jade Carving': 19, 'Line draw': 20, 'Emoji': 21
}

lora_paths = {
    "9 frame": "/tiamat-NAS/songyiren/FYP/liucheng/makeanything_models/makeanything/asymmetric_lora/asymmetric_lora_9f_general.safetensors",
    "4 frame": "/tiamat-NAS/songyiren/FYP/liucheng/makeanything_models/makeanything/asymmetric_lora/asymmetric_lora_4f_general.safetensors"
}
BASE_FLUX_CHECKPOINT = "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/unet/flux1-dev-fp8.safetensors"
# LORA_WEIGHTS_PATH="/tiamat-NAS/songyiren/FYP/liucheng/makeanything_models/makeanything/asymmetric_lora/asymmetric_lora_9f_general.safetensors"
CLIP_L_PATH = "/tiamat-NAS/hailong/storage_backup/models/stabilityai/stable-diffusion-3-medium/text_encoders/clip_l.safetensors"
T5XXL_PATH = "/tiamat-NAS/songyiren/FYP/liucheng/ComfyUI/models/clip/t5xxl_fp8_e4m3fn.safetensors"
AE_PATH = "/tiamat-vePFS/share_data/storage/huggingface/models/black-forest-labs/FLUX.1-dev/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 with dynamic paths based on the selected model
def load_target_model(frame, domain):
    global model, clip_l, t5xxl, ae, lora_model

    logger.info("Loading models...")
    # try:
    if model is None is None or clip_l is None or t5xxl is None or ae is None:
        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.")
    # Load models
    _, model = flux_utils.load_flow_model(
        BASE_FLUX_CHECKPOINT, torch.float8_e4m3fn, "cpu", disable_mmap=False
    )
    # Load LoRA weights
    LORA_WEIGHTS_PATH = lora_paths[frame]
    multiplier = 1.0
    weights_sd = load_file(LORA_WEIGHTS_PATH)
    lora_ups_num = 10 if frame=="9 frame" else 21
    lora_model, _ = lora_flux.create_network_from_weights(multiplier, None, ae, [clip_l, t5xxl], model, weights_sd, True, lora_ups_num=lora_ups_num)
    for sub_lora in lora_model.unet_loras:
        sub_lora.set_lora_up_cur(domain_index[domain]-1)

    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()

    logger.info("Models loaded successfully.")
    return "Models loaded successfully. Using Frame: {}, Damain: {}".format(frame, domain)

    # except Exception as e:
    #     logger.error(f"Error loading models: {e}")
    #     return f"Error loading models: {e}"

# The function to generate image from a prompt and conditional image
# @spaces.GPU(duration=180)
def infer(prompt, frame, 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

    frame_num = int(frame[0:1])

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

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

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

    # 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

    packed_latent_height, packed_latent_width = math.ceil(height / 16), math.ceil(width / 16)
    
    torch.manual_seed(seed)
    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)

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

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

    model.to(device)

    print(f"Model device: {model.device}")
    print(f"Noise device: {noise.device}")
    print(f"Image IDs device: {img_ids.device}")
    print(f"T5 output device: {t5_out.device}")
    print(f"Text IDs device: {txt_ids.device}")
    print(f"L pooled device: {l_pooled.device}")

    # Run the denoising process
    with accelerator.autocast(), torch.no_grad():
        x = flux_train_utils.denoise(
                model,
                noise,
                img_ids,
                t5_out,
                txt_ids,
                l_pooled,
                timesteps,
                guidance=4.0,
                t5_attn_mask=t5_attn_mask,
                cfg_scale=1.0,
            )
    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

def update_domains(floor):
    domains_dict = {
        "4 frame": [
            "LEGO", "Cook", "Painting", "Icon", "Landscape illustration", 
            "Portrait", "Transformer", "Sand art", "Illustration", "Sketch",
            "Clay toys", "Clay sculpture", "Zbrush Modeling", "Wood sculpture", "Ink painting",
            "Pencil sketch", "Fabric toys", "Oil painting", "Jade Carving", "Line draw", "Emoji"
        ],
        "9 frame": [
            "LEGO", "Cook", "Painting", "Icon", "Landscape illustration", 
            "Portrait", "Transformer", "Sand art", "Illustration", "Sketch"
        ]
    }
    return gr.Dropdown.update(choices=domains_dict[floor], label="Select Domains")

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Asymmertric LoRA Generation")

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                with gr.Column(scale=1):
                    frame_selector = gr.Radio(choices=["4 frame", "9 frame"], label="Select Floor")
                with gr.Column(scale=2):
                    domain_selector = gr.Dropdown(choices=[], label="Select Domains")
                   
            # Load Model Button
            load_button = gr.Button("Load Model")

        with gr.Column(scale=1):
            # Status message box
            status_box = gr.Textbox(label="Status", placeholder="Model loading status", interactive=False, value="Model not loaded", lines=3)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input for the prompt
            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here", lines=8)
            with gr.Row():
                seed = gr.Slider(0, np.iinfo(np.int32).max, step=1, label="Seed", value=42)
                run_button = gr.Button("Generate Image")
        
        with gr.Column(scale=1):
            # Output result
            result_image = gr.Image(label="Generated Image", interactive=False)

    frame_selector.change(update_domains, inputs=frame_selector, outputs=domain_selector)

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

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

    # gr.Markdown("### Examples")
    # examples = [
    #     [
    #         "sks14, 2*2 puzzle of 4 sub-images, step-by-step wood sculpture carving process",               # prompt
    #         "./gradio_examples/wood_sculpture.png",
    #         "Wood Sculpture",                                             # recraft_model
    #         12345                                                         # seed
    #     ],
    #     [
    #         "sks1, 3*3 puzzle of 9 sub-images, step-by-step lego model construction process",              # prompt
    #         "./gradio_examples/lego.png", 
    #         "LEGO",                                                       # recraft_model
    #         42                                                            # seed
    #     ],
    #     [
    #         "sks6, 3*3 puzzle of 9 sub-images, step-by-step portrait painting process",            # prompt
    #         "./gradio_examples/portrait.png", 
    #         "Portrait",                                                     # recraft_model
    #         999                                                           # seed
    #     ],
    #     [
    #         "sks10, 3*3 puzzle of 9 sub-images, step-by-step sketch painting process,",              # prompt
    #         "./gradio_examples/sketch.png", 
    #         "Sketch",
    #         2023
    #     ]
    # ]

    # gr.Examples(
    #     examples=examples,
    #     inputs=[prompt, sample_image, recraft_model, seed],
    #     outputs=[result_image],
    #     cache_examples=False
    # )

    # Launch the Gradio app
    demo.launch(server_port=8289, server_name="0.0.0.0", share=True)