from __future__ import annotations
import spaces
import random
import yaml
import gradio as gr
import numpy as np
import PIL
import torch
from PIL import Image, ImageDraw
from diffusers import EDMEulerScheduler
from gradio_imageslider import ImageSlider
from models import load_models
from prompt_engineer import enhance_prompt
from utils import set_timesteps_patched
import time

# Load configuration
with open("config.yaml", "r") as f:
    config = yaml.safe_load(f)

# Load models and apply patch
pipe_fast, refiner, pipe_edit, pipe_outpaint = load_models(config)
EDMEulerScheduler.set_timesteps = set_timesteps_patched


def can_expand(source_width, source_height, target_width, target_height, alignment):
    if alignment in ("Left", "Right") and source_width >= target_width:
        return False
    if alignment in ("Top", "Bottom") and source_height >= target_height:
        return False
    return True


@spaces.GPU
def infer(
    image, width, height, overlap_width, num_inference_steps, prompt_input=None, alignment="Middle"
):
    source = image
    target_size = (width, height)
    overlap = overlap_width

    # Upscale if source is smaller than target in both dimensions
    if source.width < target_size[0] and source.height < target_size[1]:
        scale_factor = min(
            target_size[0] / source.width, target_size[1] / source.height
        )
        new_width = int(source.width * scale_factor)
        new_height = int(source.height * scale_factor)
        source = source.resize((new_width, new_height), Image.LANCZOS)

    if source.width > target_size[0] or source.height > target_size[1]:
        scale_factor = min(
            target_size[0] / source.width, target_size[1] / source.height
        )
        new_width = int(source.width * scale_factor)
        new_height = int(source.height * scale_factor)
        source = source.resize((new_width, new_height), Image.LANCZOS)

    if not can_expand(
        source.width, source.height, target_size[0], target_size[1], alignment
    ):
        alignment = "Middle"

    # Calculate margins based on alignment
    if alignment == "Middle":
        margin_x = (target_size[0] - source.width) // 2
        margin_y = (target_size[1] - source.height) // 2
    elif alignment == "Left":
        margin_x = 0
        margin_y = (target_size[1] - source.height) // 2
    elif alignment == "Right":
        margin_x = target_size[0] - source.width
        margin_y = (target_size[1] - source.height) // 2
    elif alignment == "Top":
        margin_x = (target_size[0] - source.width) // 2
        margin_y = 0
    elif alignment == "Bottom":
        margin_x = (target_size[0] - source.width) // 2
        margin_y = target_size[1] - source.height

    background = Image.new("RGB", target_size, (255, 255, 255))
    background.paste(source, (margin_x, margin_y))

    mask = Image.new("L", target_size, 255)
    mask_draw = ImageDraw.Draw(mask)

    # Adjust mask generation based on alignment
    if alignment == "Middle":
        mask_draw.rectangle(
            [
                (margin_x + overlap, margin_y + overlap),
                (margin_x + source.width - overlap, margin_y + source.height - overlap),
            ],
            fill=0,
        )
    elif alignment == "Left":
        mask_draw.rectangle(
            [
                (margin_x, margin_y),
                (margin_x + source.width - overlap, margin_y + source.height),
            ],
            fill=0,
        )
    elif alignment == "Right":
        mask_draw.rectangle(
            [
                (margin_x + overlap, margin_y),
                (margin_x + source.width, margin_y + source.height),
            ],
            fill=0,
        )
    elif alignment == "Top":
        mask_draw.rectangle(
            [
                (margin_x, margin_y),
                (margin_x + source.width, margin_y + source.height - overlap),
            ],
            fill=0,
        )
    elif alignment == "Bottom":
        mask_draw.rectangle(
            [
                (margin_x, margin_y + overlap),
                (margin_x + source.width, margin_y + source.height),
            ],
            fill=0,
        )

    cnet_image = background.copy()
    cnet_image.paste(0, (0, 0), mask)

    final_prompt = f"{prompt_input} , high quality, 4k"

    (
        prompt_embeds,
        negative_prompt_embeds,
        pooled_prompt_embeds,
        negative_pooled_prompt_embeds,
    ) = pipe_outpaint.encode_prompt(final_prompt, "cuda", True)

    image = pipe_outpaint(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        image=cnet_image,
        num_inference_steps=num_inference_steps,
    )
    return image


# Generator function (refactored for better task handling and error handling)
@spaces.GPU(duration=60, queue=False)
def generate_image(
    task_type,
    input_image_edit,
    instruction,
    negative_prompt="",
    steps=25,
    randomize_seed=True,
    seed=2404,
    width=1024,
    height=1024,
    guidance_scale=3,
    fast=True,
    overlap_width=42,
    alignment="Middle",
    progress=gr.Progress(track_tqdm=True),
):
    
    torch.cuda.empty_cache()
    
    if task_type == "Image Editing" or task_type == "Outpainting":
        try:
            input_img = Image.open(input_image_edit).convert("RGB")
        except (FileNotFoundError, PIL.UnidentifiedImageError) as e:
            print(f"Error loading image: {e}")
            return None, f"Error: {e}", gr.update()  # No update for history
    else:
        input_img = None  # No input image needed for generation

    if task_type == "Image Editing":
        if randomize_seed:
            seed = random.randint(0, 999999)
        generator = torch.manual_seed(seed)
        
        output_image = pipe_edit(
            instruction,
            negative_prompt=negative_prompt,
            image=input_img,
            guidance_scale=guidance_scale*2,
            image_guidance_scale=1.5,
            width=input_img.width,
            height=input_img.height,
            num_inference_steps=steps,
            generator=generator,
            output_type="latent",
        ).images
        
        refine = refiner(
            prompt=f"{instruction}, 4k, hd, high quality, masterpiece",
            negative_prompt=negative_prompt,
            guidance_scale=7.5,
            num_inference_steps=steps,
            image=output_image,
            generator=generator,
        ).images[0]

        image_history.append(refine)  # Add output image to history

        return seed, refine, gr.update(value=image_history) 

    elif task_type == "Image Generation":
        try:
            if randomize_seed:
                seed = random.randint(0, 999999)
            generator = torch.Generator().manual_seed(seed)
            if fast:
                refine = pipe_fast(
                    prompt=instruction,
                    guidance_scale=guidance_scale,
                    num_inference_steps=int(steps / 3),
                    width=width,
                    height=height,
                    generator=generator,
                ).images[0]
            else:
                image = pipe_fast(
                    prompt=instruction,
                    negative_prompt=negative_prompt,
                    guidance_scale=guidance_scale,
                    num_inference_steps=steps,
                    width=width,
                    height=height,
                    generator=generator,
                    output_type="latent",
                ).images

                refine = refiner(
                    prompt=f"{instruction}, 4k, hd, high quality, masterpiece",
                    negative_prompt=negative_prompt,
                    guidance_scale=7.5,
                    num_inference_steps=steps,
                    image=image,
                    generator=generator,
                ).images[0]

            image_history.append(refine)  # Add output image to history (no input for generation)

            return seed, refine, gr.update(value=image_history)
        except Exception as e:  # You might want to catch more specific exceptions
            print(f"Error during image generation: {e}")
            return None, f"Error: {e}", gr.update() # No update for history

    elif task_type == "Outpainting":
        try:

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

            outpainted_image = infer(
                input_img,
                width,
                height,
                overlap_width,
                int(steps/2),
                instruction,
                alignment,
            )

            image_history.append(outpainted_image) 

            return seed, outpainted_image, gr.update(value=image_history)

        except Exception as e:
            print(f"Error during outpainting: {e}")
            return None, f"Error: {e}", gr.update() # No update for history

    else:
        raise ValueError(f"Invalid task type: {task_type}")


# Gradio UI (separated for better organization)
def create_ui(refiner, pipe_fast, pipe_edit):
    css = """
    .gradio-container {
        /* max-width: 700px !important; */ 
        h1 {text-align: center}
        footer {visibility: hidden}
    }
    """

    examples = [
        [
            "Image Generation",
            None,
            "A luxurious supercar with a unique design. The car should have a pearl white finish, and gold accents. 4k, realistic.",
        ],
        ["Image Editing", "./supercar.png", "make it red"],
        ["Image Editing", "./red_car.png", "add some snow"],
        [
            "Image Generation",
            None,
            "An alien grasping a sign board containing the word 'ALIEN' with Neon Glow, neon, futuristic, neonpunk, neon lights",
        ],
        ["Image Generation", None, "Beautiful Eiffel Tower at Night"],
        ["Image Generation", None, "Beautiful Eiffel Tower at Night"],
        ["Outpainting", "./example_1.webp", "Extend the image to the right"],
    ]

    with gr.Blocks(css=css) as demo:
        gr.Markdown("# Image Generation, Image Editing, Outpainting\n### Note: First image generation takes time")
                
        with gr.Row():
            instruction = gr.Textbox(lines=1, label="Instruction", interactive=True)
            task_type = gr.Dropdown(
                    ["Image Generation", "Image Editing", "Outpainting"],
                    label="Task",
                    value="Image Generation",
                    interactive=True,
                    scale=0
                )
            fast = gr.Checkbox(label="FAST Generation", value=True, scale=0, visible=True)
            generate_button = gr.Button("Run", scale=0)
            
        with gr.Row():
            input_image_edit = gr.Image(label="Input Image", type="filepath", interactive=True, visible=False)
            with gr.Column():
                input_image = gr.Image(label="Output Image", type="filepath", interactive=True)
                use_as_input_button = gr.Button("Use this image as Input")
                
        with gr.Row():
            guidance_scale = gr.Number( value=3.0, step=0.1, label="Guidance Scale", interactive=True )
            steps = gr.Number(value=25, step=1, label="Steps", interactive=True)
            overlap_width = gr.Slider(
                label="Overlap Width", minimum=1, maximum=100, step=1, value=42, visible=False
            )
            alignment = gr.Dropdown(
                ["Middle", "Left", "Right", "Top", "Bottom"],
                label="Alignment",
                value="Middle",
                visible=False
            )
            width = gr.Slider(
                label="Width", minimum=512, maximum=2048, step=16, value=1024, visible=True
            )
            height = gr.Slider(
                label="Height", minimum=512, maximum=2048, step=16, value=1024, visible=True
            )

        with gr.Accordion("Advanced options", open=False):
            with gr.Row():
                negative_prompt = gr.Text(
                    label="Negative prompt",
                    max_lines=1,
                    value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, ugly, disgusting, blurry, amputation,(face asymmetry, eyes asymmetry, deformed eyes, open mouth)",
                    visible=True,
                )
            with gr.Row():
                randomize_seed = gr.Checkbox(
                    label="Randomize Seed", value=True, interactive=True
                )
                seed = gr.Number(value=2404, step=1, label="Seed", interactive=True)

        gr.Examples(
            examples=examples,
            inputs=[task_type, input_image_edit, instruction],
            fn=generate_image,
            outputs=[input_image],  
            cache_examples=False,
        )

        global image_history
        image_history = []
        history_gallery = gr.Gallery(label="History", show_label=True, columns=6)

        def use_as_input(output_image):
            return gr.update(value=output_image)

        use_as_input_button.click(
            fn=use_as_input, 
            inputs=[input_image], 
            outputs=[input_image_edit] 
        )

        def update_visibility(task_type):
            if task_type == "Image Generation":
                return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)  # Hide alignment and overlap_width
            elif task_type == "Image Editing":
                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)  # Hide alignment, overlap_width, width, and height 
            elif task_type == "Outpainting":
                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True), gr.update(visible=True) # Show alignment, overlap_width, width, and height


        task_type.change(
            fn=update_visibility,
            inputs=task_type,
            outputs=[input_image_edit, fast, guidance_scale, width, height, alignment, overlap_width],  
        )

        generate_button.click(
            fn=generate_image,
            inputs=[
                task_type,
                input_image_edit,  
                instruction,
                negative_prompt,
                steps,
                randomize_seed,
                seed,
                width,
                height,
                guidance_scale,
                fast,
                overlap_width,
                alignment,
            ],
            outputs=[seed, input_image, history_gallery]
        )

        instruction.submit(
            fn=generate_image,
            inputs=[
                task_type,
                input_image_edit,  
                instruction,
                negative_prompt,
                steps,
                randomize_seed,
                seed,
                width,
                height,
                guidance_scale,
                fast,
                overlap_width,
                alignment,
            ],
            outputs=[seed, input_image, history_gallery]
        )

    return demo


# Launch the Gradio interface
demo = create_ui(refiner, pipe_fast, pipe_edit)
demo.queue(max_size=500).launch()




Add image editing v2 (Inpainting) to it. Here is the code. Below.

import gradio as gr
import spaces
import torch
from diffusers import AutoencoderKL, TCDScheduler
from diffusers.models.model_loading_utils import load_state_dict
from gradio_imageslider import ImageSlider
from huggingface_hub import hf_hub_download

from controlnet_union import ControlNetModel_Union
from pipeline_fill_sd_xl import StableDiffusionXLFillPipeline

MODELS = {
    "RealVisXL V5.0 Lightning": "SG161222/RealVisXL_V5.0_Lightning",
}

config_file = hf_hub_download(
    "xinsir/controlnet-union-sdxl-1.0",
    filename="config_promax.json",
)

config = ControlNetModel_Union.load_config(config_file)
controlnet_model = ControlNetModel_Union.from_config(config)
model_file = hf_hub_download(
    "xinsir/controlnet-union-sdxl-1.0",
    filename="diffusion_pytorch_model_promax.safetensors",
)
state_dict = load_state_dict(model_file)
model, _, _, _, _ = ControlNetModel_Union._load_pretrained_model(
    controlnet_model, state_dict, model_file, "xinsir/controlnet-union-sdxl-1.0"
)
model.to(device="cuda", dtype=torch.float16)

vae = AutoencoderKL.from_pretrained(
    "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
).to("cuda")

pipe = StableDiffusionXLFillPipeline.from_pretrained(
    "SG161222/RealVisXL_V5.0_Lightning",
    torch_dtype=torch.float16,
    vae=vae,
    controlnet=model,
    variant="fp16",
).to("cuda")

pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)


@spaces.GPU
def fill_image(prompt, image, model_selection):
    (
        prompt_embeds,
        negative_prompt_embeds,
        pooled_prompt_embeds,
        negative_pooled_prompt_embeds,
    ) = pipe.encode_prompt(prompt, "cuda", True)

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

    alpha_channel = mask.split()[3]
    binary_mask = alpha_channel.point(lambda p: p > 0 and 255)
    cnet_image = source.copy()
    cnet_image.paste(0, (0, 0), binary_mask)

    for image in pipe(
        prompt_embeds=prompt_embeds,
        negative_prompt_embeds=negative_prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
        image=cnet_image,
    ):
        yield image, cnet_image

    image = image.convert("RGBA")
    cnet_image.paste(image, (0, 0), binary_mask)

    yield source, cnet_image


def clear_result():
    return gr.update(value=None)


title = """<h1 align="center">Diffusers Fast Inpaint</h1>
<div align="center">Draw the mask over the subject you want to erase or change and write what you want to inpaint it with.</div>
<div align="center">This is a lighting model with almost no CFG and 12 steps, so don't expect high quality generations.</div>
"""

with gr.Blocks(fill_height=True, fill_width=True) as demo:
    gr.HTML(title)
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(
                label="Prompt",
                info="Describe what to inpaint the mask with",
                lines=3,
            )
        with gr.Column():
            model_selection = gr.Dropdown(
                choices=list(MODELS.keys()),
                value="RealVisXL V5.0 Lightning",
                label="Model",
            )
            run_button = gr.Button("Generate")

    with gr.Row():
        input_image = gr.ImageMask(
            type="pil",
            label="Input Image",
            crop_size=(1024, 1024),
            layers=False
        )

        result = ImageSlider(
            interactive=False,
            label="Generated Image",
        )

    use_as_input_button = gr.Button("Use as Input Image", visible=False)  

    def use_output_as_input(output_image):
        return gr.update(value=output_image[1])

    use_as_input_button.click(
        fn=use_output_as_input,
        inputs=[result],
        outputs=[input_image]
    )

    run_button.click(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=False),
        inputs=None,
        outputs=use_as_input_button,
    ).then(
        fn=fill_image,
        inputs=[prompt, input_image, model_selection],
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=use_as_input_button,
    )

    prompt.submit(
        fn=clear_result,
        inputs=None,
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=False),
        inputs=None,
        outputs=use_as_input_button,
    ).then(
        fn=fill_image,
        inputs=[prompt, input_image, model_selection],
        outputs=result,
    ).then(
        fn=lambda: gr.update(visible=True),
        inputs=None,
        outputs=use_as_input_button,
    )


demo.launch(share=False)



It uses same model and configs as outpainting so do not load it again.

Add image editing v2 (Inpainting) to it. Here is the code. Below.