FLUX.Dev-LORA-Serverless

Running on Zero

File size: 8,833 Bytes

import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import  DiffusionPipeline, FlowMatchEulerDiscreteScheduler, AutoencoderTiny, AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer,T5EncoderModel, T5TokenizerFast
from live_preview_helpers import calculate_shift, retrieve_timesteps, flux_pipe_call_that_returns_an_iterable_of_images


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

taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype).to(device)
good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=dtype).to(device)
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=dtype, vae=taef1).to(device)
torch.cuda.empty_cache()

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

pipe.flux_pipe_call_that_returns_an_iterable_of_images = flux_pipe_call_that_returns_an_iterable_of_images.__get__(pipe)

article_text = """
<div style="text-align: center;">
    <p>Enjoying the tool? Buy me a coffee and get exclusive prompt guides!</p>
    <p><i>Instantly unlock helpful tips for creating better prompts!</i></p>
    <div style="display: flex; justify-content: center;">
        <a href="https://piczify.lemonsqueezy.com/buy/0f5206fa-68e8-42f6-9ca8-4f80c587c83e">
            <img src="https://www.buymeacoffee.com/assets/img/custom_images/yellow_img.png" 
                 alt="Buy Me a Coffee" 
                 style="height: 40px; width: auto; box-shadow: 0 4px 8px rgba(0, 0, 0, 0.2); border-radius: 10px;">
        </a>
    </div>
</div>
"""

@spaces.GPU()
def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, guidance_scale=3.5, num_inference_steps=28, lora_id=None, lora_scale=0.95, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    
    # for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
    #         prompt=prompt,
    #         guidance_scale=guidance_scale,
    #         num_inference_steps=num_inference_steps,
    #         width=width,
    #         height=height,
    #         generator=generator,
    #         output_type="pil",
    #         good_vae=good_vae,
    #     ):
    #         yield img, seed

    # Handle LoRA loading
    # Load LoRA weights and prepare joint_attention_kwargs
    if lora_id:
        pipe.unload_lora_weights()
        pipe.load_lora_weights(lora_id)
        joint_attention_kwargs = {"scale": lora_scale}
    else:
        joint_attention_kwargs = None
    
    try:
        # Call the custom pipeline function with the correct keyword argument
        for img in pipe.flux_pipe_call_that_returns_an_iterable_of_images(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator,
            output_type="pil",
            good_vae=good_vae,  # Assuming good_vae is defined elsewhere
            joint_attention_kwargs=joint_attention_kwargs,  # Fixed parameter name
        ):
            yield img, seed
    finally:
        # Unload LoRA weights if they were loaded
        if lora_id:
            pipe.unload_lora_weights()

# def query(lora_id, prompt, steps=28, cfg_scale=3.5, randomize_seed=True, seed=-1, width=1024, height=1024):
#     if prompt == "" or prompt == None:
#         return None

#     if lora_id.strip() == "" or lora_id == None:
#         lora_id = "black-forest-labs/FLUX.1-dev" 

#     key = random.randint(0, 999)

#     API_URL = "https://api-inference.huggingface.co/models/"+ lora_id.strip()
    
#     API_TOKEN = random.choice([os.getenv("HF_READ_TOKEN")])
#     headers = {"Authorization": f"Bearer {API_TOKEN}"}
    
#     # prompt = GoogleTranslator(source='ru', target='en').translate(prompt)
#     # print(f'\033[1mGeneration {key} translation:\033[0m {prompt}')

#     prompt = f"{prompt} | ultra detail, ultra elaboration, ultra quality, perfect."
#     # print(f'\033[1mGeneration {key}:\033[0m {prompt}')

#     # If seed is -1, generate a random seed and use it
#     if randomize_seed:
#         seed = random.randint(1, 4294967296)
    
#     payload = {
#         "inputs": prompt,
#         "steps": steps,
#         "cfg_scale": cfg_scale,
#         "seed": seed,
#         "parameters": {
#             "width": width,  # Pass the width to the API
#             "height": height  # Pass the height to the API
#         }
#     }

#     response = requests.post(API_URL, headers=headers, json=payload, timeout=timeout)
#     if response.status_code != 200:
#         print(f"Error: Failed to get image. Response status: {response.status_code}")
#         print(f"Response content: {response.text}")
#         if response.status_code == 503:
#             raise gr.Error(f"{response.status_code} : The model is being loaded")
#         raise gr.Error(f"{response.status_code}")
    
#     try:
#         image_bytes = response.content
#         image = Image.open(io.BytesIO(image_bytes))
#         print(f'\033[1mGeneration {key} completed!\033[0m ({prompt})')
#         return image, seed, seed
#     except Exception as e:
#         print(f"Error when trying to open the image: {e}")
#         return None


examples = [
    "a tiny astronaut hatching from an egg on the moon",
    "a cat holding a sign that says hello world",
    "an anime illustration of a wiener schnitzel",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 960px;
}
.generate-btn {
    background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
    border: none !important;
    color: white !important;
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
}
"""

with gr.Blocks(css=css) as app:
    gr.HTML("<center><h1>FLUX.1-Dev with LoRA support</h1></center>")
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column():
                with gr.Row():
                    text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=3, elem_id="prompt-text-input")
                with gr.Row():
                    custom_lora = gr.Textbox(label="Custom LoRA", info="LoRA Hugging Face path (optional)", placeholder="multimodalart/vintage-ads-flux")
                with gr.Row():
                    with gr.Accordion("Advanced Settings", open=False):
                        with gr.Row():
                            lora_scale = gr.Slider(
                                label="LoRA Scale",
                                minimum=0,
                                maximum=2,
                                step=0.01,
                                value=0.95,
                            )
                            width = gr.Slider(label="Width", value=1024, minimum=64, maximum=1216, step=8)
                            height = gr.Slider(label="Height", value=1024, minimum=64, maximum=1216, step=8)
                        seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=4294967296, step=1)
                        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                        with gr.Row():
                            steps = gr.Slider(label="Inference steps steps", value=28, minimum=1, maximum=100, step=1)
                            cfg = gr.Slider(label="Guidance Scale", value=3.5, minimum=1, maximum=20, step=0.5)
                        # method = gr.Radio(label="Sampling method", value="DPM++ 2M Karras", choices=["DPM++ 2M Karras", "DPM++ SDE Karras", "Euler", "Euler a", "Heun", "DDIM"])

                with gr.Row():
                    # text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
                    text_button = gr.Button("✨ Generate Image", variant='primary', elem_classes=["generate-btn"])
            with gr.Column():
                with gr.Row():
                    image_output = gr.Image(type="pil", label="Image Output", elem_id="gallery")
                with gr.Row():
                    seed_output = gr.Textbox(label="Seed Used", show_copy_button = True)
        
        gr.Markdown(article_text)
        with gr.Column():
            gr.Examples(
                examples = examples,
                inputs = [text_prompt],
            )

        
        # text_button.click(query, inputs=[custom_lora, text_prompt, steps, cfg, randomize_seed, seed, width, height], outputs=[image_output,seed_output, seed])
        text_button.click(infer, inputs=[text_prompt, seed, randomize_seed, width, height, cfg, steps, custom_lora, lora_scale], outputs=[image_output,seed_output, seed])

app.launch()