app.py

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

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