# PyTorch 2.8 (temporary hack) import os os.system('pip install --upgrade --pre --extra-index-url https://download.pytorch.org/whl/nightly/cu126 "torch<2.9" spaces') # Actual demo code import gradio as gr import numpy as np import spaces import torch import random from PIL import Image, ImageOps from diffusers import FluxKontextPipeline from diffusers.utils import load_image # from optimization import optimize_pipeline_ MAX_SEED = np.iinfo(np.int32).max pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda") pipe.load_lora_weights("ovi054/Draw2Photo-Kontext-LoRA") pipe.fuse_lora() # optimize_pipeline_(pipe, image=Image.new("RGB", (512, 512)), prompt='prompt') import os EXAMPLES_DIR = "examples" BASE_EXAMPLES = [os.path.join(EXAMPLES_DIR, "base", f) for f in sorted(os.listdir(os.path.join(EXAMPLES_DIR, "base")))] FACE_EXAMPLES = [os.path.join(EXAMPLES_DIR, "face", f) for f in sorted(os.listdir(os.path.join(EXAMPLES_DIR, "face")))] def add_overlay(base_img, overlay_img, margin=20): """ Pastes an overlay image onto the top-right corner of a base image. The overlay is resized to be 1/5th of the width of the base image, maintaining its aspect ratio. Args: base_img (PIL.Image.Image): The main image. overlay_img (PIL.Image.Image): The image to place on top. margin (int, optional): The pixel margin from the top and right edges. Defaults to 20. Returns: PIL.Image.Image: The combined image. """ if base_img is None or overlay_img is None: return base_img base = base_img.convert("RGBA") overlay = overlay_img.convert("RGBA") # --- MODIFICATION --- # Calculate the target width to be 1/5th of the base image's width target_width = base.width // 5 # Keep aspect ratio, resize overlay to the newly calculated target width w, h = overlay.size # Add a check to prevent division by zero if the overlay image has no width if w == 0: return base new_height = int(h * (target_width / w)) overlay = overlay.resize((target_width, new_height), Image.LANCZOS) # Position: top-right corner with a margin x = base.width - overlay.width - margin y = margin # Paste the resized overlay onto the base image using its alpha channel for transparency base.paste(overlay, (x, y), overlay) return base @spaces.GPU(duration=45) def infer(input_image, input_image_upload, overlay_image, prompt="make it real", seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)): """ 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 (dict or PIL.Image.Image): The input from the gr.Paint component. input_image_upload (PIL.Image.Image): The input from the gr.Image upload component. overlay_image (PIL.Image.Image): The face photo to overlay. prompt (str): Text description of the desired edit to apply to the image. seed (int, optional): Random seed for reproducible generation. randomize_seed (bool, optional): If True, generates a random seed. guidance_scale (float, optional): Controls how closely the model follows the prompt. steps (int, optional): Controls how many steps to run the diffusion model for. progress (gr.Progress, optional): Gradio progress tracker. Returns: tuple: A 4-tuple containing the result image, the processed input image, the seed, and a gr.Button update. """ if randomize_seed: seed = random.randint(0, MAX_SEED) # --- CORRECTED LOGIC STARTS HERE --- # 1. Prioritize the uploaded image. If it exists, it becomes our main 'input_image'. if input_image_upload is not None: processed_input_image = input_image_upload # 2. If no image was uploaded, check the drawing canvas. elif isinstance(input_image, dict): # Extract the actual image from the dictionary provided by gr.Paint if "composite" in input_image and input_image["composite"] is not None: processed_input_image = input_image["composite"] elif "background" in input_image and input_image["background"] is not None: processed_input_image = input_image["background"] else: # The canvas is empty, so there's no input image. processed_input_image = None else: # Fallback in case the input is neither from upload nor a valid canvas dict. processed_input_image = None # --- CORRECTED LOGIC ENDS HERE --- # From this point on, 'processed_input_image' is either a PIL Image or None. if processed_input_image is not None: if overlay_image is not None: # Now this function is guaranteed to receive a PIL Image. processed_input_image = add_overlay(processed_input_image, overlay_image) processed_input_image = processed_input_image.convert("RGB") image = pipe( image=processed_input_image, prompt=prompt, guidance_scale=guidance_scale, width = processed_input_image.size[0], height = processed_input_image.size[1], num_inference_steps=steps, generator=torch.Generator().manual_seed(seed), ).images[0] else: # Handle the text-to-image case where no input image was provided. image = pipe( prompt=prompt, guidance_scale=guidance_scale, num_inference_steps=steps, generator=torch.Generator().manual_seed(seed), ).images[0] return image, processed_input_image, seed, gr.Button(visible=False) @spaces.GPU def infer_example(input_image, prompt): image, seed, _ = infer(input_image, prompt) return image, seed # css=""" # #col-container { # margin: 0 auto; # max-width: 960px; # } # """ css="" with gr.Blocks(css=css) as demo: with gr.Column(elem_id="col-container"): gr.Markdown(f"""# FLUX.1 Kontext [dev] + Draw2Photo LoRA Turn drawing+face into a realistic photo with FLUX.1 Kontext [dev] + [Draw2Photo LoRA](https://huggingface.co/ovi054/Draw2Photo-Kontext-LoRA) """) with gr.Row(): with gr.Column(): gr.Markdown("Step 1. Select/Upload/Draw a person ⬇️") # input_image = gr.Image(label="Upload drawing", type="pil") with gr.Row(): with gr.Tabs() as tabs: with gr.TabItem("Upload"): input_image_upload = gr.Image(label="Upload drawing", type="pil") with gr.TabItem("Draw"): input_image = gr.Paint( type="pil", brush=gr.Brush(default_size=6, colors=["#000000"], color_mode="fixed"), canvas_size = (1200,1200), layers = False ) gr.Examples( examples=[[img] for img in BASE_EXAMPLES], inputs=[input_image_upload], ) with gr.Column(): gr.Markdown("Step 2. Select/Upload a face photo ⬇️") with gr.Row(): overlay_image = gr.Image(label="Upload face photo", type="pil") gr.Examples( examples=[[img] for img in FACE_EXAMPLES], inputs=[overlay_image], ) with gr.Column(): gr.Markdown("Step 3. Press “Run” to get results ⬇️") with gr.Row(): run_button = gr.Button("Run") with gr.Accordion("Advanced Settings", open=False): prompt = gr.Text( label="Prompt", max_lines=1, value = "make it real", placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')", container=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 ) result = gr.Image(label="Result", show_label=False, interactive=False) result_input = 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_upload, prompt], # outputs=[result, seed], # fn=infer_example, # cache_examples="lazy" # ) gr.on( triggers=[run_button.click, prompt.submit], fn = infer, inputs = [input_image, input_image_upload, overlay_image, prompt, seed, randomize_seed, guidance_scale, steps], outputs = [result, result_input, seed, reuse_button] ) # reuse_button.click( # fn = lambda image: image, # inputs = [result], # outputs = [input_image] # ) demo.launch(mcp_server=True)