frogleo's picture
Update app.py
f30cd13 verified
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image
from diffusers import FluxKontextPipeline
from diffusers.utils import load_image
MAX_SEED = np.iinfo(np.int32).max
pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
@spaces.GPU
def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress()):
"""
Perform image editing using the FLUX.1 Kontext pipeline.
This function takes an input image and a text prompt to generate a modified version
of the image based on the provided instructions. It uses the FLUX.1 Kontext model
for contextual image editing tasks.
Args:
input_image (PIL.Image.Image): The input image to be edited. Will be converted
to RGB format if not already in that format.
prompt (str): Text description of the desired edit to apply to the image.
Examples: "Remove glasses", "Add a hat", "Change background to beach".
seed (int, optional): Random seed for reproducible generation. Defaults to 42.
Must be between 0 and MAX_SEED (2^31 - 1).
randomize_seed (bool, optional): If True, generates a random seed instead of
using the provided seed value. Defaults to False.
guidance_scale (float, optional): Controls how closely the model follows the
prompt. Higher values mean stronger adherence to the prompt but may reduce
image quality. Range: 1.0-10.0. Defaults to 2.5.
steps (int, optional): Controls how many steps to run the diffusion model for.
Range: 1-30. Defaults to 28.
progress (gr.Progress, optional): Gradio progress tracker for monitoring
generation progress. Defaults to gr.Progress(track_tqdm=True).
Returns:
tuple: A 3-tuple containing:
- PIL.Image.Image: The generated/edited image
- int: The seed value used for generation (useful when randomize_seed=True)
- gr.update: Gradio update object to make the reuse button visible
Example:
>>> edited_image, used_seed, button_update = infer(
... input_image=my_image,
... prompt="Add sunglasses",
... seed=123,
... randomize_seed=False,
... guidance_scale=2.5
... )
"""
progress(0,desc="Starting")
def callback_fn(pipe, step, timestep, callback_kwargs):
print(f"[Step {step}] Timestep: {timestep}")
progress_value = (step+1.0)/steps
progress(progress_value, desc=f"Image generating, {step + 1}/{steps} steps")
return callback_kwargs
if randomize_seed:
seed = random.randint(0, MAX_SEED)
if input_image:
input_image = input_image.convert("RGB")
image = pipe(
image=input_image,
prompt=prompt,
guidance_scale=guidance_scale,
width = input_image.size[0],
height = input_image.size[1],
num_inference_steps=steps,
callback_on_step_end=callback_fn,
generator=torch.Generator().manual_seed(seed),
).images[0]
else:
image = pipe(
prompt=prompt,
guidance_scale=guidance_scale,
num_inference_steps=steps,
callback_on_step_end=callback_fn,
generator=torch.Generator().manual_seed(seed),
).images[0]
progress(1, desc="Complete")
return image, seed, gr.Button(visible=True)
@spaces.GPU
def infer_example(input_image, prompt):
image, seed, _ = infer(input_image, prompt)
return image, seed
title = "# Image to Image AI Editor"
description = "Your Image-to-Image AI editor. Just describe changes (‘brighter, remove object, cartoon style’) and let the AI handle the rest—no Photoshop skills needed. Try the stable version at [Image to Image AI Generator](https://www.image2image.ai)."
with gr.Blocks() as demo:
gr.Markdown(title)
gr.Markdown(description)
with gr.Column():
with gr.Row():
with gr.Column():
input_image = gr.Image(label="Upload the image for editing", type="pil")
with gr.Row():
prompt = gr.Text(
label="Prompt",
show_label=False,
max_lines=1,
placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
container=False,
)
run_button = gr.Button("Run", scale=0)
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)
guidance_scale = gr.Slider(
label="Guidance Scale",
minimum=1,
maximum=10,
step=0.1,
value=2.5,
)
steps = gr.Slider(
label="Steps",
minimum=1,
maximum=30,
value=28,
step=1
)
with gr.Column():
result = gr.Image(label="Result", show_label=False, interactive=False)
reuse_button = gr.Button("Reuse this image", visible=False)
examples = gr.Examples(
examples=[
["flowers.png", "turn the flowers into sunflowers"],
["monster.png", "make this monster ride a skateboard on the beach"],
["cat.png", "make this cat happy"]
],
inputs=[input_image, prompt],
outputs=[result, seed],
fn=infer_example,
cache_examples="lazy"
)
gr.on(
triggers=[run_button.click, prompt.submit],
fn = infer,
inputs = [input_image, prompt, seed, randomize_seed, guidance_scale, steps],
outputs = [result, seed, reuse_button]
)
reuse_button.click(
fn = lambda image: image,
inputs = [result],
outputs = [input_image]
)
demo.launch(mcp_server=True)