app.py

import os
import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, AutoencoderKL
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from share_btn import community_icon_html, loading_icon_html, share_js
from cog_sdxl_dataset_and_utils import TokenEmbeddingsHandler
import lora
import copy
import json
import gc
# import random

import inspect
from gradio import routes
from typing import List, Type

import base64
from io import BytesIO
from PIL import Image

MY_TOKEN = os.environ.get("MY_TOKEN")

print(torch.cuda.is_available())

def image_to_base64(image: Image.Image) -> str:
    buffered = BytesIO()
    image.save(buffered, format="PNG")
    img_str = base64.b64encode(buffered.getvalue()).decode()
    return img_str


def get_types(cls_set: List[Type], component: str):
    docset = []
    types = []
    if component == "input":
        for cls in cls_set:
            doc = inspect.getdoc(cls)
            doc_lines = doc.split("\n")
            docset.append(doc_lines[1].split(":")[-1])
            types.append(doc_lines[1].split(")")[0].split("(")[-1])
    else:
        for cls in cls_set:
            doc = inspect.getdoc(cls)
            doc_lines = doc.split("\n")
            docset.append(doc_lines[-1].split(":")[-1])
            types.append(doc_lines[-1].split(")")[0].split("(")[-1])
    return docset, types
routes.get_types = get_types



with open("sdxl_loras.json", "r") as file:
    data = json.load(file)
    sdxl_loras_raw = [
        {
            "image": item["image"],
            "title": item["title"],
            "repo": item["repo"],
            "trigger_word": item["trigger_word"],
            "weights": item["weights"],
            "is_compatible": item["is_compatible"],
            "is_pivotal": item.get("is_pivotal", False),
            "text_embedding_weights": item.get("text_embedding_weights", None),
            "likes": item.get("likes", 0),
            "downloads": item.get("downloads", 0),
            "is_nc": item.get("is_nc", False)
        }
        for item in data
    ]

device = "cuda"

state_dicts = {}

for item in sdxl_loras_raw:
    saved_name = hf_hub_download(item["repo"], item["weights"])

    if not saved_name.endswith('.safetensors'):
        state_dict = torch.load(saved_name)
        # state_dict = torch.load(saved_name, map_location=torch.device('cpu'))

    else:
        state_dict = load_file(saved_name)

    state_dicts[item["repo"]] = {
        "saved_name": saved_name,
        "state_dict": state_dict
    }

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
)

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    vae=vae,
    torch_dtype=torch.float16,
)

original_pipe = copy.deepcopy(pipe)
pipe.to(device)

last_lora = ""
last_merged = False
last_fused = False

def update_selection(selected_state: gr.SelectData, sdxl_loras):
    lora_repo = sdxl_loras[selected_state.index]["repo"]
    instance_prompt = sdxl_loras[selected_state.index]["trigger_word"]
    new_placeholder = "Type a prompt. This LoRA applies for all prompts, no need for a trigger word" if instance_prompt == "" else "Type a prompt to use your selected LoRA"
    weight_name = sdxl_loras[selected_state.index]["weights"]
    updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨ {'(non-commercial LoRA, `cc-by-nc`)' if sdxl_loras[selected_state.index]['is_nc'] else '' }"
    
    return (
        updated_text,
        instance_prompt,
        gr.update(placeholder=new_placeholder),
        selected_state,
    )




def check_selected(selected_state):
    if not selected_state:
        raise gr.Error("You must select a LoRA dude")


def merge_incompatible_lora(full_path_lora, lora_scale):
    for weights_file in [full_path_lora]:
        if ";" in weights_file:
            weights_file, multiplier = weights_file.split(";")
            multiplier = float(multiplier)
        else:
            multiplier = lora_scale

        lora_model, weights_sd = lora.create_network_from_weights(
            multiplier,
            full_path_lora,
            pipe.vae,
            pipe.text_encoder,
            pipe.unet,
            for_inference=True,
        )
        lora_model.merge_to(
            pipe.text_encoder, pipe.unet, weights_sd, torch.float16, "cuda"
        )
        del weights_sd
        del lora_model
        gc.collect()


def run_lora(prompt, negative, lora_scale, selected_state, sdxl_loras, progress=gr.Progress(track_tqdm=True)):
    global last_lora, last_merged, last_fused, pipe

    print("✅ Running LoRAAAAA  >>>>>>>>>>>>> >>>>>>>>>>>>>>>>>>>")
    # print("prompt: ", prompt)
    # print("negative: ", negative)
    # print("lora_scale: ", lora_scale)
    # print("selected_state: ", selected_state)
    print("selected_state index: ", selected_state.index)


    if negative == "":
        negative = None

    if not selected_state:
        raise gr.Error("You must select a LoRA")
    repo_name = sdxl_loras[selected_state.index]["repo"]
    weight_name = sdxl_loras[selected_state.index]["weights"]

    full_path_lora = state_dicts[repo_name]["saved_name"]
    loaded_state_dict = state_dicts[repo_name]["state_dict"]
    cross_attention_kwargs = None
    if last_lora != repo_name:
        if last_merged:
            del pipe
            gc.collect()
            pipe = copy.deepcopy(original_pipe)
            pipe.to(device)
        elif (last_fused):
            pipe.unfuse_lora()
            pipe.unload_lora_weights()
        is_compatible = sdxl_loras[selected_state.index]["is_compatible"]

        if is_compatible:
            pipe.load_lora_weights(loaded_state_dict)
            pipe.fuse_lora(lora_scale)
            last_fused = True
        else:
            is_pivotal = sdxl_loras[selected_state.index]["is_pivotal"]
            if (is_pivotal):
                pipe.load_lora_weights(loaded_state_dict)
                pipe.fuse_lora(lora_scale)
                last_fused = True

                # Add the textual inversion embeddings from pivotal tuning models
                text_embedding_name = sdxl_loras[selected_state.index]["text_embedding_weights"]
                text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
                tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
                embedding_path = hf_hub_download(
                    repo_id=repo_name, filename=text_embedding_name, repo_type="model")
                embhandler = TokenEmbeddingsHandler(text_encoders, tokenizers)
                embhandler.load_embeddings(embedding_path)

            else:
                merge_incompatible_lora(full_path_lora, lora_scale)
                last_fused = False
            last_merged = True

    image = pipe(
        prompt=prompt,
        negative_prompt=negative,
        width=512,
        height=512,
        num_inference_steps=20,
        guidance_scale=7.5,
    ).images[0]
    last_lora = repo_name
    gc.collect()
    print("✅ Returning image  >>>>>>>>>>>>> >>>>>>>>>>>>>>>>>>>")
    print("image: ", image)

    return image, gr.update(visible=True)


def run_lora_light(prompt, negative, lora_scale, selected_index, sdxl_loras, progress=gr.Progress(track_tqdm=True)):
    global last_lora, last_merged, last_fused, pipe

    print("✅ Running run_lora_light  >>>>>>>>>>>>> >>>>>>>>>>>>>>>>>>>")
    print("prompt: ", prompt)
    print("negative: ", negative)
    print("lora_scale: ", lora_scale)
    print("selected_state: ", selected_index)

    if negative == "":
        negative = None

    # if not selected_state:
    #     raise gr.Error("You must select a LoRA")

    repo_name = sdxl_loras[selected_index]["repo"]
    weight_name = sdxl_loras[selected_index]["weights"]

    full_path_lora = state_dicts[repo_name]["saved_name"]
    loaded_state_dict = state_dicts[repo_name]["state_dict"]
    cross_attention_kwargs = None
    if last_lora != repo_name:
        if last_merged:
            del pipe
            gc.collect()
            pipe = copy.deepcopy(original_pipe)
            pipe.to(device)
        elif (last_fused):
            pipe.unfuse_lora()
            pipe.unload_lora_weights()
        is_compatible = sdxl_loras[selected_index]["is_compatible"]

        if is_compatible:
            pipe.load_lora_weights(loaded_state_dict)
            pipe.fuse_lora(lora_scale)
            last_fused = True
        else:
            is_pivotal = sdxl_loras[selected_index]["is_pivotal"]
            if (is_pivotal):
                pipe.load_lora_weights(loaded_state_dict)
                pipe.fuse_lora(lora_scale)
                last_fused = True

                # Add the textual inversion embeddings from pivotal tuning models
                text_embedding_name = sdxl_loras[selected_index]["text_embedding_weights"]
                text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
                tokenizers = [pipe.tokenizer, pipe.tokenizer_2]
                embedding_path = hf_hub_download(
                    repo_id=repo_name, filename=text_embedding_name, repo_type="model")
                embhandler = TokenEmbeddingsHandler(text_encoders, tokenizers)
                embhandler.load_embeddings(embedding_path)

            else:
                merge_incompatible_lora(full_path_lora, lora_scale)
                last_fused = False
            last_merged = True

    image = pipe(
        prompt=prompt,
        negative_prompt=negative,
        width=512,
        height=512,
        num_inference_steps=20,
        guidance_scale=7.5,
    ).images[0]
    last_lora = repo_name
    gc.collect()
    print("image: ", image)

    image_base64 = image_to_base64(image)

    return image_base64



def shuffle_gallery(sdxl_loras):
    order = "likes"
    sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get(order, 0), reverse=True)
    return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery

def swap_gallery(order, sdxl_loras):
    if(order == "random"):
        return shuffle_gallery(sdxl_loras)
    else:
        sorted_gallery = sorted(sdxl_loras, key=lambda x: x.get(order, 0), reverse=True)
        return [(item["image"], item["title"]) for item in sorted_gallery], sorted_gallery


# App code
def hallo(full_string):
    print("✅ Hallo >>>>>>>>>")
    print("string: ", full_string)

    parts = full_string.split("+")
    text_part = parts[0].strip()
    number_part = parts[1].strip()
    idx_lora = int(number_part)
    token_part = parts[2].strip()

    if(token_part == MY_TOKEN):
        img_result = run_lora_light(prompt=text_part, negative="No naked bodies", lora_scale=0.8, selected_index=idx_lora, sdxl_loras=sdxl_loras_raw)
        return img_result
    else:
        img_result = {"message": "Failed request", "prompt": text_part}
        return img_result

    

def hadet(x):
    return f"Hadet, {x}"



with gr.Blocks(css="custom.css") as demo:
    gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
    # <<<<<<< new additions
    t = gr.Textbox()
    b = gr.Button("Hallo")
    a = gr.Button("Hadet")
    o = gr.Textbox()
    b.click(hallo, inputs=[t], outputs=[o])
    a.click(hadet, inputs=[t], outputs=[o])
    #  new additions >>>>>>>>>

    title = gr.HTML(
        """<h1>Algorithmic Dream Interpreter | Art Generator</h1>""",
        elem_id="title",
    )
    selected_state = gr.State()
    print("✅ selected_state: ", selected_state)
    with gr.Row():
        with gr.Box(elem_id="gallery_box"):
            order_gallery = gr.Radio(choices=[
                                     "random", "likes"], value="random", label="Order by", elem_id="order_radio")
            gallery = gr.Gallery(
                #value=[(item["image"], item["title"]) for item in sdxl_loras],
                label="SDXL LoRA Gallery",
                allow_preview=False,
                columns=4,
                elem_id="gallery",
                show_share_button=False,
                # height=784
                height=384
            )

        with gr.Column():
            prompt_title = gr.Markdown(
                value="### Click on a LoRA in the gallery to select it",
                visible=True,
                elem_id="selected_lora",
            )
            with gr.Row():
                prompt = gr.Textbox(label="Prompt", show_label=False, lines=1, max_lines=1,
                                    placeholder="Type a prompt after selecting a LoRA", elem_id="prompt")
                button = gr.Button("Run", elem_id="run_button")
            with gr.Group(elem_id="share-btn-container", visible=False) as share_group:
                community_icon = gr.HTML(community_icon_html)
                loading_icon = gr.HTML(loading_icon_html)
                share_button = gr.Button(
                    "Share to community", elem_id="share-btn")
            result = gr.Image(
                interactive=False, label="Generated Image", elem_id="result-image"
            )
            with gr.Accordion("Advanced options", open=False):
                negative = gr.Textbox(label="Negative Prompt")
                weight = gr.Slider(
                    0, 10, value=0.8, step=0.1, label="LoRA weight")
                
    gallery.select(
        fn=update_selection,
        inputs=[gr_sdxl_loras],

        outputs=[prompt_title, prompt, prompt,
                 selected_state],
        queue=False,
        show_progress=False
    )

    
    order_gallery.change(
        fn=swap_gallery,
        inputs=[order_gallery, gr_sdxl_loras],
        outputs=[gallery, gr_sdxl_loras],
        queue=False
    )

    
    button.click(
        fn=check_selected,
        inputs=[selected_state],
        queue=False,
        show_progress=False
    ).success(
        fn=run_lora,
        inputs=[prompt, negative, weight, selected_state, gr_sdxl_loras],
        outputs=[result, share_group],
    )

    # share_button.click(None, [], [], _js=share_js)

    demo.load(fn=shuffle_gallery, inputs=[gr_sdxl_loras], outputs=[
              gallery, gr_sdxl_loras], queue=False)
    

# <<<<<<< new additions
ifa = gr.Interface(lambda: None, inputs=[t], outputs=[o])
demo.input_components = ifa.input_components
demo.output_components = ifa.output_components
demo.examples = None
demo.predict_durations = [] 
# new additions >>>>>>>>>

# demo.queue(max_size=20)
demo.launch()