Create app.py

app.py

@@ -0,0 +1,234 @@
+from typing import Tuple, Optional
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+import random
+from PIL import Image
+import json
+import boto3
+from io import BytesIO
+from datetime import datetime
+from huggingface_hub import login
+import os
+
+from diffusers import FluxKontextPipeline
+from diffusers.utils import load_image
+from diffusers.utils import load_image, make_image_grid
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+login(token=HF_TOKEN)
+
+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: Image.Image, 
+    prompt: str, 
+    seed: int, 
+    randomize_seed: bool, 
+    guidance_scale: float, 
+    steps: int, 
+    progress
+):
+    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,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+    else:
+        image = pipe(
+            prompt=prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=steps,
+            generator=torch.Generator().manual_seed(seed),
+        ).images[0]
+    return image
+
+
+def process(
+    image_url: str, 
+    prompt: str, 
+    seed:int, 
+    randomize_seed:bool, 
+    guidance_scale:float, 
+    steps:int, 
+    upload_to_r2:bool, 
+    account_id: str, 
+    access_key: str, 
+    secret_key: str, 
+    bucket:str, 
+    progress=gr.Progress(track_tqdm=True)
+):
+    result = {"status": "false", "message": ""}     
+    input_image = load_image(image_url)
+    if not isinstance(input_image, Image.Image):
+        result["status"] = "fail"
+        result["message"] = "Invalid input image url"
+        return [], json.dumps(result)
+    
+    try:
+        generated_image = infer(input_image, prompt, seed, randomize_seed, guidance_scale, steps, progress)
+    except Exception as e:
+        result["status"] = "faield"
+        result["message"] = "generate image failed"
+        generated_image = None
+    
+    if generated_image:
+        if upload_to_r2:
+            url = upload_image_to_r2(generated_image, account_id, access_key, secret_key, bucket)
+            result = {"status": "success", "message": "upload image success", "url": url}    
+        else:
+            result = {"status": "success", "message": "Image generated but not uploaded"}
+    final_images = []
+    
+    if isinstance(input_image, Image.Image):
+        final_images.append(input_image)
+    if isinstance(generated_image, Image.Image):
+        final_images.append(generated_image)
+    
+    progress(100, "finish!")
+    return final_images, json.dumps(result)
+
+
+def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
+    with calculateDuration("Upload image"):
+        print("upload_image_to_r2", account_id, access_key, secret_key, bucket_name)
+        connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
+        s3 = boto3.client(
+            's3',
+            endpoint_url=connectionUrl,
+            region_name='auto',
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key
+        )
+        current_time = datetime.now().strftime("%Y/%m/%d/%H%M%S")
+        image_file = f"generated_images/{current_time}_{random.randint(0, MAX_SEED)}.png"
+        buffer = BytesIO()
+        image.save(buffer, "PNG")
+        buffer.seek(0)
+        s3.upload_fileobj(buffer, bucket_name, image_file)
+        print("upload finish", image_file)
+
+        # start to generate thumbnail
+        thumbnail = image.copy()
+        thumbnail_width = 256
+        aspect_ratio = image.height / image.width
+        thumbnail_height = int(thumbnail_width * aspect_ratio)
+        thumbnail = thumbnail.resize((thumbnail_width, thumbnail_height), Image.LANCZOS)
+        
+        # Generate the thumbnail image filename
+        thumbnail_file = image_file.replace(".png", "_thumbnail.png")
+        
+        # Save thumbnail to buffer and upload
+        thumbnail_buffer = BytesIO()
+        thumbnail.save(thumbnail_buffer, "PNG")
+        thumbnail_buffer.seek(0)
+        s3.upload_fileobj(thumbnail_buffer, bucket_name, thumbnail_file)
+        print("upload thumbnail finish", thumbnail_file)
+        
+    return image_file
+    
+
+with gr.Blocks() as demo:
+    
+    with gr.Column():
+        gr.Markdown(f"""# FLUX.1 Kontext [dev]
+Image editing and manipulation model guidance-distilled from FLUX.1 Kontext [pro], [[blog]](https://bfl.ai/announcements/flux-1-kontext-dev) [[model]](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev)
+        """)
+        with gr.Row():
+            with gr.Column():
+                image_url =  gr.Text(
+                    label="Orginal image url",
+                    show_label=True,
+                    max_lines=1,
+                    placeholder="Enter image url for inpainting",
+                    container=False
+                )
+                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(value='Submit', variant='primary', 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.Accordion("R2 Settings", open=False):
+                    upload_to_r2 = gr.Checkbox(label="Upload to R2", value=False)
+                    with gr.Row():
+                        account_id = gr.Textbox(label="Account Id", placeholder="Enter R2 account id")
+                        bucket = gr.Textbox(label="Bucket Name", placeholder="Enter R2 bucket name here")
+    
+                    with gr.Row():
+                        access_key = gr.Textbox(label="Access Key", placeholder="Enter R2 access key here")
+                        secret_key = gr.Textbox(label="Secret Key", placeholder="Enter R2 secret key here")
+                
+                    
+            with gr.Column():
+                generated_images = gr.Gallery(label="Result", show_label=True)
+                output_json_component = gr.Code(label="JSON Result", language="json")
+
+    run_button.click(
+        fn = process,
+        inputs = [
+            image_url, 
+            prompt, 
+            seed, 
+            randomize_seed, 
+            guidance_scale, 
+            steps, 
+            upload_to_r2, 
+            account_id, 
+            access_key, 
+            secret_key, 
+            bucket
+        ],
+        outputs = [
+            generated_images, 
+            output_json_component
+        ]
+    )
+    
+
+demo.queue()
+demo.launch(share=True)