app.py

import spaces

import gradio as gr
import numpy as np
import os
import random
import json
from PIL import Image
import torch
from torchvision import transforms
import zipfile

from diffusers import FluxFillPipeline, AutoencoderKL
from PIL import Image
# from samgeo.text_sam import LangSAM

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

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# sam = LangSAM(model_type="sam2-hiera-large").to(device)

pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")

with open("lora_models.json", "r") as f:
    lora_models = json.load(f)

def download_model(model_name, model_path):
    print(f"Downloading model: {model_name} from {model_path}")
    try:
        pipe.load_lora_weights(model_path)
        print(f"Successfully downloaded model: {model_name}")
    except Exception as e:
        print(f"Failed to download model: {model_name}. Error: {e}")

# Iterate through the models and download each one
for model_name, model_path in lora_models.items():
    download_model(model_name, model_path)

lora_models["None"] = None

# def calculate_optimal_dimensions(image: Image.Image):
#     # Extract the original dimensions
#     original_width, original_height = image.size

#     # Set constants
#     MIN_ASPECT_RATIO = 9 / 16
#     MAX_ASPECT_RATIO = 16 / 9
#     FIXED_DIMENSION = 1024

#     # Calculate the aspect ratio of the original image
#     original_aspect_ratio = original_width / original_height

#     # Determine which dimension to fix
#     if original_aspect_ratio > 1:  # Wider than tall
#         width = FIXED_DIMENSION
#         height = round(FIXED_DIMENSION / original_aspect_ratio)
#     else:  # Taller than wide
#         height = FIXED_DIMENSION
#         width = round(FIXED_DIMENSION * original_aspect_ratio)

#     # Ensure dimensions are multiples of 8
#     width = (width // 8) * 8
#     height = (height // 8) * 8

#     # Enforce aspect ratio limits
#     calculated_aspect_ratio = width / height
#     if calculated_aspect_ratio > MAX_ASPECT_RATIO:
#         width = (height * MAX_ASPECT_RATIO // 8) * 8
#     elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
#         height = (width / MIN_ASPECT_RATIO // 8) * 8

#     # Ensure width and height remain above the minimum dimensions
#     width = max(width, 576) if width == FIXED_DIMENSION else width
#     height = max(height, 576) if height == FIXED_DIMENSION else height

#     return width, height

@spaces.GPU(durations=300)
def infer(edit_images, prompt, width, height, lora_model, strength, seed=42, randomize_seed=False, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
    # pipe.enable_xformers_memory_efficient_attention()
    gr.Info("Infering")

    if lora_model != "None":
        pipe.load_lora_weights(lora_models[lora_model])
        pipe.enable_lora()

    gr.Info("starting checks")

    image = edit_images["background"]
    mask = edit_images["layers"][0]

    if not image:
        gr.Info("Please upload an image.")
        return None, None


    # width, height = calculate_optimal_dimensions(image)
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # controlImage = processor(image)
    gr.Info("generating image")
    image = pipe(
        # mask_image_latent=vae.encode(controlImage),
        prompt=prompt,
        prompt_2=prompt,
        image=image,
        mask_image=mask,
        height=height,
        width=width,
        guidance_scale=guidance_scale,
        # strength=strength,
        num_inference_steps=num_inference_steps,
        generator=torch.Generator(device='cuda').manual_seed(seed),
        # generator=torch.Generator().manual_seed(seed),
        # lora_scale=0.75 // not supported in this version
    ).images[0]

    output_image_jpg = image.convert("RGB")
    output_image_jpg.save("output.jpg", "JPEG")

    return output_image_jpg, seed
    # return image, seed

def download_image(image):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image.save("output.png", "PNG")
    return "output.png"

def save_details(result, edit_image, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps):
    image = edit_image["background"]
    mask = edit_image["layers"][0]

    if isinstance(result, np.ndarray):
        result = Image.fromarray(result)
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    if isinstance(mask, np.ndarray):
        mask = Image.fromarray(mask)

    result.save("saved_result.png", "PNG")
    image.save("saved_image.png", "PNG")
    mask.save("saved_mask.png", "PNG")

    details = {
        "prompt": prompt,
        "lora_model": lora_model,
        "strength": strength,
        "seed": seed,
        "guidance_scale": guidance_scale,
        "num_inference_steps": num_inference_steps
    }

    with open("details.json", "w") as f:
        json.dump(details, f)

    # Create a ZIP file
    with zipfile.ZipFile("output.zip", "w") as zipf:
        zipf.write("saved_result.png")
        zipf.write("saved_image.png")
        zipf.write("saved_mask.png")
        zipf.write("details.json")

    return "output.zip"

def set_image_as_inpaint(image):
    return image

# def generate_mask(image, click_x, click_y):
#     text_prompt = "face"
#     mask = sam.predict(image, text_prompt, box_threshold=0.24, text_threshold=0.24)
#     return mask

examples = [
    "photography of a young woman,  accent lighting,  (front view:1.4),  "
    # "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: 1000px;
}
"""

with gr.Blocks(css=css) as demo:

    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""# FLUX.1 [dev]
        """)
        with gr.Row():
            with gr.Column():
                edit_image = gr.ImageEditor(
                    label='Upload and draw mask for inpainting',
                    type='pil',
                    sources=["upload", "webcam"],
                    image_mode='RGB',
                    layers=False,
                    brush=gr.Brush(colors=["#FFFFFF"]),
                    # height=600
                )
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=2,
                    placeholder="Enter your prompt",
                    container=False,
                )

                lora_model = gr.Dropdown(
                    label="Select LoRA Model",
                    choices=list(lora_models.keys()),
                    value="None",
                )

                run_button = gr.Button("Run")

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

        with gr.Accordion("Advanced Settings", open=False):

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

                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1,
                    maximum=30,
                    step=0.5,
                    value=50,
                )

                num_inference_steps = gr.Slider(
                    label="Number of inference steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=28,
                )

            with gr.Row():

                strength = gr.Slider(
                    label="Strength",
                    minimum=0,
                    maximum=1,
                    step=0.01,
                    value=0.85,
                )

            with gr.Row():

                width = gr.Slider(
                    label="width",
                    minimum=512,
                    maximum=3072,
                    step=1,
                    value=1024,
                )

                height = gr.Slider(
                    label="height",
                    minimum=512,
                    maximum=3072,
                    step=1,
                    value=1024,
                )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn = infer,
        inputs = [edit_image, prompt, width, height, lora_model, strength, seed, randomize_seed, guidance_scale, num_inference_steps],
        outputs = [result, seed]
    )

    download_button = gr.Button("Download Image as PNG")
    set_inpaint_button = gr.Button("Set Image as Inpaint")
    save_button = gr.Button("Save Details")

    download_button.click(
            fn=download_image,
            inputs=[result],
            outputs=gr.File(label="Download Image")
        )

    set_inpaint_button.click(
            fn=set_image_as_inpaint,
            inputs=[result],
            outputs=[edit_image]
    )

    save_button.click(
            fn=save_details,
            inputs=[result, edit_image, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps],
            outputs=gr.File(label="Download/Save Status")
    )

    # edit_image.select(
    #     fn=generate_mask,
    #     inputs=[edit_image, gr.Number(), gr.Number()],
    #     outputs=[edit_image]
    # )

# demo.launch()
PASSWORD = os.getenv("GRADIO_PASSWORD")
USERNAME = os.getenv("GRADIO_USERNAME")
# Create an authentication object
def authenticate(username, password):
    if username == USERNAME and password == PASSWORD:
        return True

    else:
        return False
# Launch the app with authentication

demo.launch(share=True, debug=True, auth=authenticate)