app.py

import json
import random
import requests
import os
from PIL import Image

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import DiffusionPipeline, LCMScheduler

def get_image(image_data):
    if isinstance(image_data, str):
        return image_data
    
    if isinstance(image_data, dict):
        local_path = image_data.get('local_path')
        hf_url = image_data.get('hf_url')
    else:
        print(f"Unexpected image_data format: {type(image_data)}")
        return None

    if local_path and os.path.exists(local_path):
        try:
            Image.open(local_path).verify()
            return local_path
        except Exception as e:
            print(f"Error loading local image {local_path}: {e}")

    if hf_url:
        try:
            response = requests.get(hf_url)
            if response.status_code == 200:
                img = Image.open(requests.get(hf_url, stream=True).raw)
                img.verify()
                img.save(local_path)
                return local_path
            else:
                print(f"Failed to fetch image from URL {hf_url}. Status code: {response.status_code}")
        except Exception as e:
            print(f"Error loading image from URL {hf_url}: {e}")

    print(f"Failed to load image for {image_data}")
    return None

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

sdxl_loras_raw = sorted(sdxl_loras_raw, key=lambda x: x["likes"], reverse=True)

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "stabilityai/stable-diffusion-xl-base-1.0"

pipe = DiffusionPipeline.from_pretrained(model_id, variant="fp16")

# Create LCMScheduler with default config
lcm_scheduler = LCMScheduler.from_config(pipe.scheduler.config)

# Remove the 'skip_prk_steps' if it exists in the config
if hasattr(lcm_scheduler.config, 'skip_prk_steps'):
    delattr(lcm_scheduler.config, 'skip_prk_steps')

pipe.scheduler = lcm_scheduler
pipe.to(device=DEVICE, dtype=torch.float16)

# Load Flash SDXL LoRA
flash_sdxl_id = "jasperai/flash-sdxl"
pipe.load_lora_weights(flash_sdxl_id, adapter_name="flash_lora")

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

def update_selection(selected_state: gr.SelectData, gr_sdxl_loras):
    lora_id = gr_sdxl_loras[selected_state.index]["repo"]
    trigger_word = gr_sdxl_loras[selected_state.index]["trigger_word"]
    return lora_id, trigger_word

@spaces.GPU
def infer(
    pre_prompt,
    prompt,
    seed,
    randomize_seed,
    num_inference_steps,
    negative_prompt,
    guidance_scale,
    user_lora_selector,
    user_lora_weight,
    progress=gr.Progress(track_tqdm=True),
):
    try:
        # Load the user-selected LoRA
        new_adapter_id = user_lora_selector.replace("/", "_")
        pipe.load_lora_weights(user_lora_selector, adapter_name=new_adapter_id)

        # Set adapter weights
        pipe.set_adapters(["flash_lora", new_adapter_id], adapter_weights=[1.0, user_lora_weight])
        gr.Info("LoRA setup complete")

        if randomize_seed:
            seed = random.randint(0, MAX_SEED)

        generator = torch.Generator().manual_seed(seed)

        if pre_prompt != "":
            prompt = f"{pre_prompt} {prompt}"

        # Use Flash Diffusion settings
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            guidance_scale=1.0,  # Flash Diffusion typically uses guidance_scale=1
            num_inference_steps=4,  # Flash Diffusion uses fewer steps
            generator=generator,
        ).images[0]

        return image
    except Exception as e:
        gr.Error(f"An error occurred: {str(e)}")
        return None

css = """
h1 {
    text-align: center;
    display:block;
}
p {
    text-align: justify;
    display:block;
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown(
        f"""
    # ⚡ FlashDiffusion: FlashLoRA ⚡
    This is an interactive demo of [Flash Diffusion](https://gojasper.github.io/flash-diffusion-project/) **on top of** existing LoRAs.
    
    The distillation method proposed in [Flash Diffusion: Accelerating Any Conditional Diffusion Model for Few Steps Image Generation](http://arxiv.org/abs/2406.02347) *by Clément Chadebec, Onur Tasar, Eyal Benaroche and Benjamin Aubin* from Jasper Research.
    The LoRAs can be added **without** any retraining for similar results in most cases. Feel free to tweak the parameters and use your own LoRAs by giving a look at the [Github Repo](https://github.com/gojasper/flash-diffusion)
    """
    )
    gr.Markdown(
        "If you enjoy the space, please also promote *open-source* by giving a ⭐ to our repo [![GitHub Stars](https://img.shields.io/github/stars/gojasper/flash-diffusion?style=social)](https://github.com/gojasper/flash-diffusion)"
    )

    gr_sdxl_loras = gr.State(value=sdxl_loras_raw)
    gr_lora_id = gr.State(value="")

    with gr.Row():
        with gr.Blocks():
            with gr.Column():
                user_lora_selector = gr.Textbox(
                    label="Current Selected LoRA",
                    max_lines=1,
                    interactive=False,
                )

                user_lora_weight = gr.Slider(
                    label="Selected LoRA Weight",
                    minimum=0.5,
                    maximum=3,
                    step=0.1,
                    value=1,
                )

                gallery = gr.Gallery(
                    value=[(item["image"], item["title"]) for item in sdxl_loras_raw],
                    label="SDXL LoRA Gallery",
                    allow_preview=False,
                    columns=3,
                    elem_id="gallery",
                    show_share_button=False,
                )

        with gr.Column():
            with gr.Row():
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter your prompt",
                    container=False,
                    scale=5,
                )

                run_button = gr.Button("Run", scale=1)

            result = gr.Image(label="Result", show_label=False)

            with gr.Accordion("Advanced Settings", open=False):
                pre_prompt = gr.Text(
                    label="Pre-Prompt",
                    show_label=True,
                    max_lines=1,
                    placeholder="Pre Prompt from the LoRA config",
                    container=True,
                    scale=5,
                )

                seed = gr.Slider(
                    label="Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )

                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)

                with gr.Row():
                    num_inference_steps = gr.Slider(
                        label="Number of inference steps",
                        minimum=4,
                        maximum=8,
                        step=1,
                        value=4,
                    )

                with gr.Row():
                    guidance_scale = gr.Slider(
                        label="Guidance Scale",
                        minimum=1,
                        maximum=6,
                        step=0.5,
                        value=1,
                    )

                hint_negative = gr.Markdown(
                    """💡 _Hint : Negative Prompt will only work with Guidance > 1 but the model was 
                    trained to be used with guidance = 1 (ie. without guidance).
                    Can degrade the results, use cautiously._"""
                )

                negative_prompt = gr.Text(
                    label="Negative Prompt",
                    show_label=False,
                    max_lines=1,
                    placeholder="Enter a negative Prompt",
                    container=False,
                )

    gr.on(
        [
            run_button.click,
            seed.change,
            randomize_seed.change,
            prompt.submit,
            negative_prompt.change,
            negative_prompt.submit,
            guidance_scale.change,
        ],
        fn=infer,
        inputs=[
            pre_prompt,
            prompt,
            seed,
            randomize_seed,
            num_inference_steps,
            negative_prompt,
            guidance_scale,
            user_lora_selector,
            user_lora_weight,
        ],
        outputs=[result],
    )

    gallery.select(
        fn=update_selection,
        inputs=[gr_sdxl_loras],
        outputs=[
            user_lora_selector,
            pre_prompt,
        ],
        show_progress="hidden",
    )

    gr.Markdown("**Disclaimer:**")
    gr.Markdown(
        "This demo is only for research purpose. Users are solely responsible for any content they create, and it is their obligation to ensure that it adheres to appropriate and ethical standards."
    )

demo.queue().launch()