app_image_edit.py

import os
import sys
import logging

import spaces
import torch

import gradio as gr
import tempfile
import numpy as np
from datetime import datetime
from pathlib import Path
from PIL import Image
import json
import base64
import time
import random
import gc
import math

from diffusers import QwenImageEditPlusPipeline, FlowMatchEulerDiscreteScheduler
from huggingface_hub import InferenceClient


from config_image_edit import TITLE, LOGO_HTML, CUSTOM_HEAD, HEADLINE_MD, SUBHEADLINE_MD, THEME, DEFAULT_PROMPT, PROMPT_PLACEHOLDER, DEFAULT_PROMPT_NEGATIVE, DATA_IMAGE_WINK, DATA_IMAGE_HEART

log = logging.getLogger(__name__)

# MARK: GLOBAL CONSTANTS:
# Define paths using pathlib.Path for consistency
BASE_DIR = Path(__file__).resolve().parent
RES = BASE_DIR / "_res"
ASSETS = RES / "assets"
EXAMPLES = BASE_DIR / "examples"
IMAGE_CACHE = BASE_DIR / "image_cache"

# Ensure the image cache directory exists
IMAGE_CACHE.mkdir(exist_ok=True)

# Set static paths for Gradio
gr.set_static_paths([str(IMAGE_CACHE), str(EXAMPLES), str(ASSETS), str(RES)])

MAX_SEED = np.iinfo(np.int32).max
MIN_IMAGE_SIZE = 256
MAX_IMAGE_SIZE = 2048


# ------------------
# MARK: Pipeline load
# ------------------
scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load the model pipeline
pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", scheduler=scheduler, torch_dtype=torch.bfloat16).to(device)
pipe.set_progress_bar_config(disable=None)

# Load LoRA for faster inference
pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning", weight_name="Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors")
pipe.fuse_lora()


SYSTEM_PROMPT = '''
# Edit Instruction Rewriter
You are a professional edit instruction rewriter. Your task is to generate a precise, concise, and visually achievable professional-level edit instruction based on the user-provided instruction and the image to be edited.  
Please strictly follow the rewriting rules below:
## 1. General Principles
- Keep the rewritten prompt **concise and comprehensive**. Avoid overly long sentences and unnecessary descriptive language.  
- If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.  
- Keep the main part of the original instruction unchanged, only enhancing its clarity, rationality, and visual feasibility.  
- All added objects or modifications must align with the logic and style of the scene in the input images.  
- If multiple sub-images are to be generated, describe the content of each sub-image individually.  
## 2. Task-Type Handling Rules
### 1. Add, Delete, Replace Tasks
- If the instruction is clear (already includes task type, target entity, position, quantity, attributes), preserve the original intent and only refine the grammar.  
- If the description is vague, supplement with minimal but sufficient details (category, color, size, orientation, position, etc.). For example:  
    > Original: "Add an animal"  
    > Rewritten: "Add a light-gray cat in the bottom-right corner, sitting and facing the camera"  
- Remove meaningless instructions: e.g., "Add 0 objects" should be ignored or flagged as invalid.  
- For replacement tasks, specify "Replace Y with X" and briefly describe the key visual features of X.  
### 2. Text Editing Tasks
- All text content must be enclosed in English double quotes `" "`. Keep the original language of the text, and keep the capitalization.  
- Both adding new text and replacing existing text are text replacement tasks, For example:  
    - Replace "xx" to "yy"  
    - Replace the mask / bounding box to "yy"  
    - Replace the visual object to "yy"  
- Specify text position, color, and layout only if user has required.  
- If font is specified, keep the original language of the font.  
### 3. Human Editing Tasks
- Make the smallest changes to the given user's prompt.  
- If changes to background, action, expression, camera shot, or ambient lighting are required, please list each modification individually.
- **Edits to makeup or facial features / expression must be subtle, not exaggerated, and must preserve the subject's identity consistency.**
    > Original: "Add eyebrows to the face"  
    > Rewritten: "Slightly thicken the person's eyebrows with little change, look natural."
### 4. Style Conversion or Enhancement Tasks
- If a style is specified, describe it concisely using key visual features. For example:  
    > Original: "Disco style"  
    > Rewritten: "1970s disco style: flashing lights, disco ball, mirrored walls, vibrant colors"  
- For style reference, analyze the original image and extract key characteristics (color, composition, texture, lighting, artistic style, etc.), integrating them into the instruction.  
- **Colorization tasks (including old photo restoration) must use the fixed template:**  
  "Restore and colorize the old photo."  
- Clearly specify the object to be modified. For example:  
    > Original: Modify the subject in Picture 1 to match the style of Picture 2.  
    > Rewritten: Change the girl in Picture 1 to the ink-wash style of Picture 2 — rendered in black-and-white watercolor with soft color transitions.
### 5. Material Replacement
- Clearly specify the object and the material. For example: "Change the material of the apple to papercut style."
- For text material replacement, use the fixed template:
    "Change the material of text "xxxx" to laser style"
### 6. Logo/Pattern Editing
- Material replacement should preserve the original shape and structure as much as possible. For example:
   > Original: "Convert to sapphire material"  
   > Rewritten: "Convert the main subject in the image to sapphire material, preserving similar shape and structure"
- When migrating logos/patterns to new scenes, ensure shape and structure consistency. For example:
   > Original: "Migrate the logo in the image to a new scene"  
   > Rewritten: "Migrate the logo in the image to a new scene, preserving similar shape and structure"
### 7. Multi-Image Tasks
- Rewritten prompts must clearly point out which image's element is being modified. For example:  
    > Original: "Replace the subject of picture 1 with the subject of picture 2"  
    > Rewritten: "Replace the girl of picture 1 with the boy of picture 2, keeping picture 2's background unchanged"  
- For stylization tasks, describe the reference image's style in the rewritten prompt, while preserving the visual content of the source image.  
## 3. Rationale and Logic Check
- Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" requires logical correction.
- Supplement missing critical information: e.g., if position is unspecified, choose a reasonable area based on composition (near subject, blank space, center/edge, etc.).
# Output Format Example
```json
{
   "Rewritten": "..."
}
'''

def polish_prompt_hf(original_prompt, img_list):
    
    # Ensure HF_TOKEN is set
    api_key = os.environ.get("HF_TOKEN")
    if not api_key:
        print("Warning: HF_TOKEN not set. Falling back to original prompt.")
        return original_prompt
    prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {original_prompt}\n\nRewritten Prompt:"
    system_prompt = "you are a helpful assistant, you should provide useful answers to users."
    try:
        # Initialize the client
        client = InferenceClient(
            provider="nebius",
            api_key=api_key,
        )

        # Convert list of images to base64 data URLs
        image_urls = []
        if img_list is not None:
            # Ensure img_list is actually a list
            if not isinstance(img_list, list):
                img_list = [img_list]
            
            for img in img_list:
                image_url = None
                # If img is a PIL Image
                if hasattr(img, 'save'):  # Check if it's a PIL Image
                    buffered = BytesIO()
                    img.save(buffered, format="PNG")
                    img_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
                    image_url = f"data:image/png;base64,{img_base64}"
                # If img is already a file path (string)
                elif isinstance(img, str):
                    with open(img, "rb") as image_file:
                        img_base64 = base64.b64encode(image_file.read()).decode('utf-8')
                    image_url = f"data:image/png;base64,{img_base64}"
                else:
                    print(f"Warning: Unexpected image type: {type(img)}, skipping...")
                    continue
                
                if image_url:
                    image_urls.append(image_url)

        # Build the content array with text first, then all images
        content = [
            {
                "type": "text",
                "text": prompt
            }
        ]
        
        # Add all images to the content
        for image_url in image_urls:
            content.append({
                "type": "image_url",
                "image_url": {
                    "url": image_url
                }
            })

        # Format the messages for the chat completions API
        messages = [
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": content
            }
        ]

        # Call the API
        completion = client.chat.completions.create(
            model="Qwen/Qwen2.5-VL-72B-Instruct",
            messages=messages,
        )
        
        # Parse the response
        result = completion.choices[0].message.content
        
        # Try to extract JSON if present
        if '"Rewritten"' in result:
            try:
                # Clean up the response
                result = result.replace('```json', '').replace('```', '')
                result_json = json.loads(result)
                polished_prompt = result_json.get('Rewritten', result)
            except:
                polished_prompt = result
        else:
            polished_prompt = result
            
        polished_prompt = polished_prompt.strip().replace("\n", " ")
        return polished_prompt
        
    except Exception as e:
        print(f"Error during API call to Hugging Face: {e}")
        # Fallback to original prompt if enhancement fails
        return original_prompt 


def encode_image(pil_image):
    import io
    buffered = io.BytesIO()
    pil_image.save(buffered, format="PNG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")


def use_history_as_input(evt: gr.SelectData, history):
    if history and evt.index < len(history):
        im = history[evt.index][0]
        image_path = f"image_cache/{history[evt.index][1]}"
        im.save(image_path, format="PNG")
        return image_path  # gr.update(value=history[evt.index][0])
    return  # gr.update()


@spaces.GPU(duration=300)
def infer(
    images,
    prompt,
    negative_prompt=" ",
    num_inference_steps=30,
    true_guidance_scale=2.0,
    seed=42,
    randomize_seed=False,
    rewrite_prompt=True,
    gallery=None,
    height=None,
    width=None,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):

    if true_guidance_scale <= 2:
        negative_prompt = " "
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Set up the generator for reproducibility
    generator = torch.Generator(device=device).manual_seed(seed)
    
    # Load input images into PIL Images
    pil_images = []
    if images is not None:
        for item in images:
            try:
                if isinstance(item[0], Image.Image):
                    pil_images.append(item[0].convert("RGB"))
                elif isinstance(item[0], str):
                    pil_images.append(Image.open(item[0]).convert("RGB"))
                elif hasattr(item, "name"):
                    pil_images.append(Image.open(item.name).convert("RGB"))
            except Exception:
                continue

    if height==256 and width==256:
        height, width = None, None
        
    print(f"Calling pipeline with prompt: '{prompt}'")
    print(f"Negative Prompt: '{negative_prompt}'")
    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
    
    if rewrite_prompt and len(pil_images) > 0:
        prompt = polish_prompt_hf(prompt, pil_images)
        print(f"Rewritten Prompt: {prompt}")
    

    # Generate the image
    image = pipe(
        image=pil_images if len(pil_images) > 0 else None,
        prompt=prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=num_images_per_prompt,
    ).images[0]
    
    timestamp = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
    image_path = IMAGE_CACHE / f"{timestamp}.png"
    
    try:
        image.save(image_path, format="PNG")
    except Exception as exc:
        raise gr.Error(f"Fehler beim Speichern des Bildes: {exc}")
    
    if gallery is None:
        gallery = []
    gallery.insert(0, (image, str(f"{timestamp}.png")))

    return str(image_path), seed, gallery


# ------------------
# MARK: GRADIO UI
# ------------------
with gr.Blocks(title=TITLE) as demo:
    with gr.Row(elem_classes="row-logo"):
        gr.HTML(
            f"""{LOGO_HTML}""",
            elem_classes="logo-html",
        )

    with gr.Tab(" Start"):
        with gr.Row(elem_classes="row-header"):
            gr.HTML(
                f"""<center>
                      <h2 class="headline">{HEADLINE_MD}</h2>
                      <h3 class="sub-headline">{SUBHEADLINE_MD}</h3>
                      <p class="like-request">lg Sebastian <img id="wink" src="{DATA_IMAGE_WINK}" width="20" /> gib dem Space gerne ein <img id="heart" src="{DATA_IMAGE_HEART}" width="20"/></p>
                    </center>
                        """,
                elem_classes="md-header",
            )
        with gr.Row(elem_classes="row-main"):
            with gr.Column(elem_classes="col-input") as input_column:
                # INPUT DATA COLUMN:
                with gr.Column(elem_classes="storyboard-column"):
                    with gr.Row(elem_classes="gradio-row storyboard-row", equal_height=True):
                        with gr.Column():
                            gr.Markdown("# **** Eingabe Bilder", elem_classes="markdown-label")
                            # input_image_component = gr.Image(type="pil", label="Start-Bild", show_label=True, elem_classes="storyboard-image image-1", interactive=True)
                            input_image_component = gr.Gallery(type="pil", label="Bilder", show_label=True, allow_preview=False, object_fit="cover", elem_classes="storyboard-image image-1", interactive=True)
                        gr.HTML("""<i class="fa-solid fa-arrow-right"></i>""", elem_classes="storyboard-arrow")
                        # input_image_last_frame_component = gr.Image(type="pil", label="End-Bild (Optional)", show_label=True, elem_classes="storyboard-image image-2")
                        with gr.Column():
                            gr.Markdown("# **** Prompt", elem_classes="markdown-label")
                            prompt = gr.Textbox(value=DEFAULT_PROMPT, placeholder=PROMPT_PLACEHOLDER, label="Regieanweisung", show_label=False, elem_classes="regieanweisung-textbox", lines=10, max_lines=12, interactive=True, html_attributes=gr.InputHTMLAttributes(autocorrect="off", spellcheck=False))

                    steps_slider = gr.Slider(
                        minimum=1,
                        maximum=30,
                        step=1,
                        value=8,
                        label="🚶‍➡️ Inferenzschritte",
                        info="Anzahl der Iterationen, die das Modell zur Bildgenerierung durchläuft; mehr Schritte erhöhen Detailtreue, verlangsamen aber die Berechnung.",
                        interactive=True,
                        elem_classes="einstellungen-slider einstellungen-steps",
                    )

                with gr.Column(elem_classes="profi-einstellungen-column"):
                    with gr.Accordion("⚙️ Profi-Einstellungen", elem_classes="profi-einstellungen-accordion", open=False):
                        with gr.Row(elem_classes="profi-einstellungen-inner-row"):
                            with gr.Column(elem_classes="profi-einstellungen-inner-column"):
                                gr.Markdown("# **** Negative Anweisungen", elem_classes="markdown-label")
                                negative_prompt_input = gr.Textbox(value=DEFAULT_PROMPT_NEGATIVE, label="Profi-Name", show_label=False, elem_classes="negative-prompt-textbox", lines=10, interactive=False, html_attributes=gr.InputHTMLAttributes(autocorrect="off", spellcheck=False))

                            with gr.Column(elem_classes="profi-einstellungen-inner-column"):

                                with gr.Row(elem_classes="einstellungen-seed-row"):
                                    seed_input = gr.Number(value=42, minimum=0, maximum=MAX_SEED, step=1, label="🌱 Seed", info="Ein fester Ausgangswert, der die Zufallszahlen im Bildgenerierungsprozess steuert, damit das gleiche Ergebnis bei gleicher Eingabe reproduzierbar ist.", elem_classes="einstellungen-seed", elem_id="seed_input", interactive=False)
                                    randomize_seed_checkbox = gr.Checkbox(value=True, label="Zufallsgenerator", show_label=False, elem_classes="einstellungen-random-seed", interactive=True)

                                guidance_scale_input = gr.Slider(
                                    minimum=0.0,
                                    maximum=10.0,
                                    step=0.1,
                                    value=2,
                                    label="🧐 Leitwert",
                                    info="Steuert, wie stark das Modell dem Prompt folgt; höhere Werte führen zu genauerer Bildgestaltung.",
                                    interactive=True,
                                    elem_classes="einstellungen-slider einstellungen-guidance-scale-1",
                                )

                                rewrite_prompt = gr.Checkbox(label="✨ Prompt verbessern", value=True, elem_classes="einstellungen-rewrite-prompt", interactive=True)

                game_button = gr.Button("Bearbeiten ✨", variant="primary", elem_classes="gradio-button run-button", elem_id="gameBtn", size="lg")
                generate_button = gr.Button("Bearbeiten ✨", variant="primary", elem_classes="gradio-button run-button", elem_id="runBtn", size="lg")

                with gr.Column(elem_classes="output-column", visible=False) as output_column_1:
                    # output_video = gr.Video(value=None, label="Ergebnis", show_label=False, elem_classes="result-video", visible=True, interactive=False)
                    with gr.Row(elem_classes="gradio-row", equal_height=True):
                        with gr.Column():
                            gr.Markdown("# **** Ergebnis", elem_classes="markdown-label")
                            result_image = gr.Image(label="Ergebnis", type="pil", show_label=False, buttons=["download", "fullscreen"], elem_classes="result-image", visible=True, interactive=False)
                        
                        with gr.Column():
                            gr.Markdown("# **** Verlauf", elem_classes="markdown-label")
                            result_history = gr.Gallery([], label="Ergebnis", show_label=False, buttons=["download", "fullscreen"], object_fit="cover", elem_classes="result-gallery", interactive=False, allow_preview=False, type="pil")

                # greet_btn.click(fn=greet, inputs=[], outputs=output, api_name="greet")
    with gr.Tab(" Info", elem_classes="info-tab"):
        with gr.Column(elem_classes="info-column"):
            gr.Markdown(
                """
                ## 👉 **Hinweis:** Dieser Space verwendet Hugging Face Zero GPU.
                - Jede Anfrage wird auf einem frischen GPU-Container ausgeführt
                - Das Modell ist bereits geladen und wird genutzt
                - Diese Technik ist umweltfreundlicher als konventionelle Methoden
                - Es optimiert außerdem die Kosten und Ressourcen-Nutzung auf Hugging Face
                - 100% Annonymes Datenhandling durch die sofortige abschaltungdes virtuellen Arbeitscontainers
                👍
                
                
                ## ⚡ **Wie Zero GPU im Detail funktioniert:**
                Stell dir vor, du hast einen superleistungsstarken Computer – aber nur für genau die Sekunden, in denen du ihn wirklich brauchst. Sobald du fertig bist, schaltet er sich wieder aus, damit niemand anders darauf warten muss und du nichts unnötig verbrauchst.
                Genau so funktioniert ZeroGPU: Diese Web-App „Bild → Video“ läuft normalerweise im „Schlafmodus“. Sobald du etwas generieren möchtest, wird blitzschnell ein leistungsstarker Rechner mit Grafikchip gestartet, erledigt die Arbeit – und schaltet sich danach wieder ab. Das spart Ressourcen, ist umweltfreundlicher und ermöglicht es, dass viele Nutzer fair und reibungslos auf dieselben leistungsstarken Maschinen zugreifen können.
                Das Beste? Du merkst davon fast nichts – außer vielleicht einem kurzen Ladebild, bevor dein Video beginnt zu entstehen.
                """,
            )

    randomize_seed_checkbox.change(fn=lambda x: gr.update(interactive=(not x)), inputs=[randomize_seed_checkbox], outputs=[seed_input], show_progress=False)
    guidance_scale_input.change(fn=lambda x: gr.update(interactive=(x > 2)), inputs=[guidance_scale_input], outputs=[negative_prompt_input], show_progress=False)

    result_history.select(fn=use_history_as_input, inputs=[result_history], outputs=[result_image], show_progress=False)
    
    ui_inputs = [input_image_component, prompt, negative_prompt_input, steps_slider, guidance_scale_input, seed_input, randomize_seed_checkbox, rewrite_prompt, result_history]
    generate_button.click(fn=lambda: gr.Column(visible=True), outputs=[output_column_1], show_progress=False, scroll_to_output=True).then(fn=infer, inputs=ui_inputs, outputs=[result_image, seed_input, result_history], show_progress_on=[result_image, result_history], scroll_to_output=True).then(
        fn=None, outputs=generate_button, js="() => { document.getElementById('gameBtn').classList.remove('process-running'); }"
    )

if __name__ == "__main__":
    # demo.launch(theme=THEME, head=CUSTOM_HEAD, css=CUSTOM_CSS, footer_links=["gradio", "settings"])
    demo.launch(theme=THEME, head=CUSTOM_HEAD, css_paths=["_res/_custom_image_edit.css", "_res/miniGameButton.css"], footer_links=["gradio", "settings"])