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

            #     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, 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(debug=True, auth=authenticate)
# demo.launch()


# import gradio as gr
# import numpy as np
# import torch
# import random
# from PIL import Image
# import cv2
# import spaces
# import os

# # ------------------ Inpainting Pipeline Setup ------------------ #
# from diffusers import FluxFillPipeline

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

# pipe = FluxFillPipeline.from_pretrained(
#     "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
# )
# pipe.load_lora_weights("alvdansen/flux-koda")
# pipe.enable_lora()

# 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 minimum dimensions are met
#     width = max(width, 576) if width == FIXED_DIMENSION else width
#     height = max(height, 576) if height == FIXED_DIMENSION else height

#     return width, height

# # ------------------ SAM (Transformers) Imports and Initialization ------------------ #
# from transformers import SamModel, SamProcessor

# # Load the model and processor from Hugging Face.
# sam_model = SamModel.from_pretrained("facebook/sam-vit-base")
# sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")

# @spaces.GPU(durations=300)
# def generate_mask_with_sam(image: Image.Image, mask_prompt: str):
#     """
#     Generate a segmentation mask using SAM (via Hugging Face Transformers).

#     The mask_prompt is expected to be a comma-separated string of two integers,
#     e.g. "450,600" representing an (x,y) coordinate in the image.

#     The function converts the coordinate into the proper input format for SAM and returns a binary mask.
#     """
#     if mask_prompt.strip() == "":
#         raise ValueError("No mask prompt provided.")

#     try:
#         # Parse the mask_prompt into a coordinate
#         coords = [int(x.strip()) for x in mask_prompt.split(",")]
#         if len(coords) != 2:
#             raise ValueError("Expected two comma-separated integers (x,y).")
#     except Exception as e:
#         raise ValueError("Invalid mask prompt. Please provide coordinates as 'x,y'. Error: " + str(e))

#     # The SAM processor expects a list of input points.
#     # Format the point as a list of lists; here we assume one point per image.
#     # (The Transformers SAM expects the points in [x, y] order.)
#     input_points = [coords]  # e.g. [[450,600]]
#     # Optionally, you can supply input_labels (1 for foreground, 0 for background)
#     input_labels = [1]

#     # Prepare the inputs for the SAM processor.
#     inputs = sam_processor(images=image,
#                            input_points=[input_points],
#                            input_labels=[input_labels],
#                            return_tensors="pt")

#     # Move tensors to the same device as the model.
#     device = next(sam_model.parameters()).device
#     inputs = {k: v.to(device) for k, v in inputs.items()}

#     # Forward pass through SAM.
#     with torch.no_grad():
#         outputs = sam_model(**inputs)

#     # The output contains predicted masks; we take the first mask from the first prompt.
#     # (Assuming outputs.pred_masks is of shape (batch_size, num_masks, H, W))
#     pred_masks = outputs.pred_masks  # Tensor of shape (1, num_masks, H, W)
#     mask = pred_masks[0][0].detach().cpu().numpy()

#     # Convert the mask to binary (0 or 255) using a threshold.
#     mask_bin = (mask > 0.5).astype(np.uint8) * 255
#     mask_pil = Image.fromarray(mask_bin)
#     return mask_pil

# # ------------------ Inference Function ------------------ #
# @spaces.GPU(durations=300)
# def infer(edit_images, prompt, mask_prompt,
#           seed=42, randomize_seed=False, width=1024, height=1024,
#           guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
#     # Get the base image from the "background" layer.
#     image = edit_images["background"]
#     width, height = calculate_optimal_dimensions(image)

#     # If a mask prompt is provided, use the SAM-based mask generator.
#     if mask_prompt and mask_prompt.strip() != "":
#         try:
#             mask = generate_mask_with_sam(image, mask_prompt)
#         except Exception as e:
#             raise ValueError("Error generating mask from prompt: " + str(e))
#     else:
#         # Fall back to using a manually drawn mask (from the first layer).
#         try:
#             mask = edit_images["layers"][0]
#         except (TypeError, IndexError):
#             raise ValueError("No mask provided. Please either draw a mask or supply a mask prompt.")

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

#     # Run the inpainting diffusion pipeline with the provided prompt and mask.
#     image_out = pipe(
#         prompt=prompt,
#         image=image,
#         mask_image=mask,
#         height=height,
#         width=width,
#         guidance_scale=guidance_scale,
#         num_inference_steps=num_inference_steps,
#         generator=torch.Generator(device='cuda').manual_seed(seed),
#     ).images[0]

#     output_image_jpg = image_out.convert("RGB")
#     output_image_jpg.save("output.jpg", "JPEG")
#     return output_image_jpg, seed

# # ------------------ Gradio UI ------------------ #
# 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("# FLUX.1 [dev] with SAM (Transformers) Mask Generation")
#         with gr.Row():
#             with gr.Column():
#                 # The image editor now allows you to optionally draw a mask.
#                 edit_image = gr.ImageEditor(
#                     label='Upload Image (and optionally draw a mask)',
#                     type='pil',
#                     sources=["upload", "webcam"],
#                     image_mode='RGB',
#                     layers=False,  # We will generate a mask automatically if needed.
#                     brush=gr.Brush(colors=["#FFFFFF"]),
#                 )
#                 prompt = gr.Text(
#                     label="Inpainting Prompt",
#                     show_label=False,
#                     max_lines=2,
#                     placeholder="Enter your inpainting prompt",
#                     container=False,
#                 )
#                 mask_prompt = gr.Text(
#                     label="Mask Prompt (enter a coordinate as 'x,y')",
#                     show_label=True,
#                     placeholder="E.g. 450,600",
#                     container=True,
#                 )
#                 generate_mask_btn = gr.Button("Generate Mask")
#                 mask_preview = gr.Image(label="Mask Preview", show_label=True)
#                 run_button = gr.Button("Run")
#             result = gr.Image(label="Result", show_label=False)

#         # Button to preview the generated mask.
#         def on_generate_mask(image, mask_prompt):
#             if image is None or mask_prompt.strip() == "":
#                 return None
#             mask = generate_mask_with_sam(image, mask_prompt)
#             return mask

#         generate_mask_btn.click(
#             fn=on_generate_mask,
#             inputs=[edit_image, mask_prompt],
#             outputs=[mask_preview]
#         )

#         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():
#                 width = gr.Slider(
#                     label="Width",
#                     minimum=256,
#                     maximum=MAX_IMAGE_SIZE,
#                     step=32,
#                     value=1024,
#                     visible=False
#                 )
#                 height = gr.Slider(
#                     label="Height",
#                     minimum=256,
#                     maximum=MAX_IMAGE_SIZE,
#                     step=32,
#                     value=1024,
#                     visible=False
#                 )
#             with gr.Row():
#                 guidance_scale = gr.Slider(
#                     label="Guidance Scale",
#                     minimum=1,
#                     maximum=30,
#                     step=0.5,
#                     value=3.5,
#                 )
#                 num_inference_steps = gr.Slider(
#                     label="Number of Inference Steps",
#                     minimum=1,
#                     maximum=50,
#                     step=1,
#                     value=28,
#                 )

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

# # 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(auth=authenticate)