Update app.py

This commit is contained in:
sigma
2026-01-04 21:07:30 +00:00
committed by system
parent e3a6c9c237
commit 428cf464ff

702
app.py
View File

@@ -1,340 +1,474 @@
import gradio as gr import gradio as gr
import os import numpy as np
import subprocess import random
import shutil
import json
import time
from pathlib import Path
import torch import torch
import spaces import spaces
from diffusers import DiffusionPipeline
# ========================================== from PIL import Image
# 1. SETUP & GLOBAL VARS from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline
# ========================================== # from optimization import optimize_pipeline_
# from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
# from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
# from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
DATASET_DIR = Path("./datasets") from huggingface_hub import InferenceClient
OUTPUT_DIR = Path("./output") import math
DATASET_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)
# global tracking for loras import os
# key: friendly name, value: path import base64
AVAILABLE_LORAS = {} from io import BytesIO
import json
print("loading z-image-turbo pipeline...") SYSTEM_PROMPT = '''
pipe = DiffusionPipeline.from_pretrained( # Edit Instruction Rewriter
"Tongyi-MAI/Z-Image-Turbo", 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.
torch_dtype=torch.bfloat16,
low_cpu_mem_usage=False,
)
pipe.to("cuda")
print("pipeline loaded!")
# ========================================== Please strictly follow the rewriting rules below:
# 2. TRAINING LOGIC
# ==========================================
def check_gpu(): ## 1. General Principles
if torch.cuda.is_available(): - Keep the rewritten prompt **concise and comprehensive**. Avoid overly long sentences and unnecessary descriptive language.
return f"✅ gpu available: {torch.cuda.get_device_name(0)}" - If the instruction is contradictory, vague, or unachievable, prioritize reasonable inference and correction, and supplement details when necessary.
return "⚠️ no gpu detected" - 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.
def upload_and_prepare_dataset(files, dataset_name, trigger_word): ## 2. Task-Type Handling Rules
if not files:
return "❌ upload images first", None, ""
if not dataset_name: ### 1. Add, Delete, Replace Tasks
dataset_name = f"dataset_{int(time.time())}" - 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.
dataset_path = DATASET_DIR / dataset_name ### 2. Text Editing Tasks
dataset_path.mkdir(exist_ok=True, parents=True) - 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.
image_count = 0 ### 3. Human Editing Tasks
for file in files: - Make the smallest changes to the given user's prompt.
if file.name.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.bmp')): - If changes to background, action, expression, camera shot, or ambient lighting are required, please list each modification individually.
filename = Path(file.name).name - **Edits to makeup or facial features / expression must be subtle, not exaggerated, and must preserve the subject's identity consistency.**
dest = dataset_path / filename > Original: "Add eyebrows to the face"
shutil.copy(file.name, dest) > Rewritten: "Slightly thicken the person's eyebrows with little change, look natural."
caption_file = dest.with_suffix('.txt') ### 4. Style Conversion or Enhancement Tasks
caption_text = trigger_word if trigger_word else "a photo" - If a style is specified, describe it concisely using key visual features. For example:
with open(caption_file, 'w') as f: > Original: "Disco style"
f.write(caption_text) > 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.
image_count += 1 ### 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"
if image_count == 0: ### 6. Logo/Pattern Editing
return "❌ no valid images found", None, "" - 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"
return f"✅ ready: {image_count} images in {dataset_name}", str(dataset_path), dataset_name ### 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.
# request 10 mins gpu for training ## 3. Rationale and Logic Check
@spaces.GPU(duration=200) - Resolve contradictory instructions: e.g., "Remove all trees but keep all trees" requires logical correction.
def train_lora( - Supplement missing critical information: e.g., if position is unspecified, choose a reasonable area based on composition (near subject, blank space, center/edge, etc.).
dataset_path,
project_name,
trigger_word,
steps,
learning_rate,
lora_rank,
resolution,
progress=gr.Progress()
):
if not dataset_path:
return "❌ no dataset", None
if not project_name: # Output Format Example
project_name = f"lora_{int(time.time())}" ```json
{
"Rewritten": "..."
}
'''
output_path = OUTPUT_DIR / project_name def polish_prompt_hf(original_prompt, img_list):
output_path.mkdir(exist_ok=True, parents=True) """
Rewrites the prompt using a Hugging Face InferenceClient.
# config generation Supports multiple images via img_list.
config = { """
"job": "extension", # Ensure HF_TOKEN is set
"config": { api_key = os.environ.get("inference_providers")
"name": project_name, if not api_key:
"process": [{ print("Warning: HF_TOKEN not set. Falling back to original prompt.")
"type": "sd_trainer", return original_prompt
"training_folder": str(output_path), prompt = f"{SYSTEM_PROMPT}\n\nUser Input: {original_prompt}\n\nRewritten Prompt:"
"device": "cuda:0", system_prompt = "you are a helpful assistant, you should provide useful answers to users."
"trigger_word": trigger_word or "",
"network": {
"type": "lora",
"linear": int(lora_rank),
"linear_alpha": int(lora_rank),
},
"save": {
"dtype": "float16",
"save_every": int(steps), # save only at end to save space
"max_step_saves_to_keep": 1,
},
"datasets": [{
"folder_path": dataset_path,
"caption_ext": "txt",
"caption_dropout_rate": 0.05,
"resolution": [int(resolution), int(resolution)],
}],
"train": {
"batch_size": 1,
"steps": int(steps),
"gradient_accumulation_steps": 1,
"train_unet": True,
"train_text_encoder": False,
"gradient_checkpointing": True,
"noise_scheduler": "flowmatch",
"optimizer": "adamw8bit",
"lr": float(learning_rate),
"ema_config": {"use_ema": True, "ema_decay": 0.99},
"dtype": "bf16",
},
"model": {
"name_or_path": "Tongyi-MAI/Z-Image-Base",
"is_v_pred": False,
"quantize": True,
},
}]
}
}
config_path = output_path / "config.json"
with open(config_path, 'w') as f:
json.dump(config, f, indent=2)
# install ai-toolkit
progress(0.1, desc="setting up environment...")
if not Path("./ai-toolkit").exists():
try: try:
subprocess.run(["git", "clone", "https://github.com/ostris/ai-toolkit.git"], check=True) # Initialize the client
subprocess.run(["pip", "install", "-q", "-r", "ai-toolkit/requirements.txt"], check=True) client = InferenceClient(
except Exception as e: provider="nebius",
return f"❌ setup failed: {e}", None api_key=api_key,
progress(0.2, desc="training (this takes time)...")
try:
# run training script
# explicitly passing environment to ensure cuda visibility in subprocess
env = os.environ.copy()
proc = subprocess.run(
["python", "ai-toolkit/run.py", str(config_path)],
capture_output=True,
text=True,
env=env,
timeout=3500
) )
if proc.returncode != 0: # Convert list of images to base64 data URLs
return f"❌ training crashed:\n{proc.stderr}", None 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]
# find result for img in img_list:
lora_files = list(output_path.glob("*.safetensors")) image_url = None
if lora_files: # If img is a PIL Image
lora_file = lora_files[-1] if hasattr(img, 'save'): # Check if it's a PIL Image
AVAILABLE_LORAS[project_name] = str(lora_file) 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
# update the dropdown choices dynamically if image_url:
choices = [("None", None)] + [(k, v) for k, v in AVAILABLE_LORAS.items()] image_urls.append(image_url)
return f"✅ trained: {project_name}", str(lora_file) # Build the content array with text first, then all images
content = [
{
"type": "text",
"text": prompt
}
]
return "⚠️ finished but no safetensors found", None # Add all images to the content
for image_url in image_urls:
content.append({
"type": "image_url",
"image_url": {
"url": image_url
}
})
except Exception as e: # Format the messages for the chat completions API
return f"❌ fatal error: {e}", None messages = [
{"role": "system", "content": system_prompt},
{
"role": "user",
"content": content
}
]
# ========================================== # Call the API
# 3. INFERENCE LOGIC completion = client.chat.completions.create(
# ========================================== model="Qwen/Qwen2.5-VL-72B-Instruct",
messages=messages,
)
@spaces.GPU # Parse the response
def generate_image( result = completion.choices[0].message.content
prompt,
height,
width,
steps,
seed,
randomize_seed,
lora_path,
lora_scale
):
# handle lora loading/unloading
pipe.unload_lora_weights() # clean slate
if lora_path and os.path.exists(lora_path): # Try to extract JSON if present
print(f"loading lora: {lora_path}") if '"Rewritten"' in result:
try: try:
pipe.load_lora_weights(lora_path) # Clean up the response
# manual scaling not always supported directly without fuse, result = result.replace('```json', '').replace('```', '')
# but usually applied by default. result_json = json.loads(result)
# for simplicitly we just load it. 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: except Exception as e:
print(f"lora load failed: {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")
# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
# Scheduler configuration for Lightning
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,
}
# Initialize scheduler with Lightning config
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
# Load the model pipeline
pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511",
scheduler=scheduler,
torch_dtype=dtype).to(device)
pipe.load_lora_weights(
"lightx2v/Qwen-Image-Edit-2511-Lightning",
weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors"
)
pipe.fuse_lora()
# # Apply the same optimizations from the first version
# pipe.transformer.__class__ = QwenImageTransformer2DModel
# pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
# # --- Ahead-of-time compilation ---
# optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
# --- UI Constants and Helpers ---
MAX_SEED = np.iinfo(np.int32).max
def use_output_as_input(output_images):
"""Convert output images to input format for the gallery"""
if output_images is None or len(output_images) == 0:
return []
return output_images
# --- Main Inference Function (with hardcoded negative prompt) ---
@spaces.GPU()
def infer(
images,
prompt,
seed=42,
randomize_seed=False,
true_guidance_scale=1.0,
num_inference_steps=4,
height=None,
width=None,
rewrite_prompt=True,
num_images_per_prompt=1,
progress=gr.Progress(track_tqdm=True),
):
"""
Run image-editing inference using the Qwen-Image-Edit pipeline.
Parameters:
images (list): Input images from the Gradio gallery (PIL or path-based).
prompt (str): Editing instruction (may be rewritten by LLM if enabled).
seed (int): Random seed for reproducibility.
randomize_seed (bool): If True, overrides seed with a random value.
true_guidance_scale (float): CFG scale used by Qwen-Image.
num_inference_steps (int): Number of diffusion steps.
height (int | None): Optional output height override.
width (int | None): Optional output width override.
rewrite_prompt (bool): Whether to rewrite the prompt using Qwen-2.5-VL.
num_images_per_prompt (int): Number of images to generate.
progress: Gradio progress callback.
Returns:
tuple: (generated_images, seed_used, UI_visibility_update)
"""
# Hardcode the negative prompt as requested
negative_prompt = " "
if randomize_seed: if randomize_seed:
seed = torch.randint(0, 2**32 - 1, (1,)).item() seed = random.randint(0, MAX_SEED)
generator = torch.Generator("cuda").manual_seed(int(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 = pipe(
image=pil_images if len(pil_images) > 0 else None,
prompt=prompt, prompt=prompt,
height=int(height), height=height,
width=int(width), width=width,
num_inference_steps=int(steps), negative_prompt=negative_prompt,
guidance_scale=0.0, num_inference_steps=num_inference_steps,
generator=generator, generator=generator,
).images[0] true_cfg_scale=true_guidance_scale,
num_images_per_prompt=num_images_per_prompt,
).images
return image, seed # Return images, seed, and make button visible
return image, seed, gr.update(visible=True)
def update_lora_list(): # --- Examples and UI Layout ---
"""helper to refresh dropdown""" examples = []
choices = [("None", None)] + [(k, v) for k, v in AVAILABLE_LORAS.items()]
return gr.Dropdown(choices=choices)
# ========================================== css = """
# 4. UI CONSTRUCTION #col-container {
# ========================================== margin: 0 auto;
max-width: 1024px;
}
#logo-title {
text-align: center;
}
#logo-title img {
width: 400px;
}
#edit_text{margin-top: -62px !important}
"""
custom_theme = gr.themes.Soft(primary_hue="yellow", secondary_hue="slate") with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
with gr.Blocks(theme=custom_theme, title="Z-Image ZeroGPU Trainer") as demo: gr.HTML("""
<div id="logo-title">
gr.Markdown("# ⚡ Z-Image-Turbo: Train & Test") <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_edit_logo.png" alt="Qwen-Image Edit Logo" width="400" style="display: block; margin: 0 auto;">
<h2 style="font-style: italic;color: #5b47d1;margin-top: -27px !important;margin-left: 96px">[Plus] Fast, 4-steps with LightX2V LoRA</h2>
with gr.Tabs(): </div>
""")
# TAB 1: INFERENCE gr.Markdown("""
with gr.Tab("🎨 Generate"): [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series.
This demo uses the new [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) with the [Qwen-Image-Lightning-2511](https://huggingface.co/lightx2v/Qwen-Image-Edit-2511-Lightning) LoRA for accelerated inference.
Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) to run locally with ComfyUI or diffusers.
""")
with gr.Row(): with gr.Row():
with gr.Column(): with gr.Column():
prompt_input = gr.Textbox(label="Prompt", lines=3) input_images = gr.Gallery(label="Input Images",
show_label=False,
type="pil",
interactive=True)
with gr.Column():
result = gr.Gallery(label="Result", show_label=False, type="pil", interactive=False)
# Add this button right after the result gallery - initially hidden
use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False)
with gr.Row(): with gr.Row():
lora_selector = gr.Dropdown( prompt = gr.Text(
label="Select LoRA", label="Prompt",
choices=[("None", None)], show_label=False,
placeholder="describe the edit instruction",
container=False,
)
run_button = gr.Button("Edit!", variant="primary")
with gr.Accordion("Advanced Settings", open=False):
# Negative prompt UI element is removed here
seed = gr.Slider(
label="Seed",
minimum=0,
maximum=MAX_SEED,
step=1,
value=0,
)
randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
with gr.Row():
true_guidance_scale = gr.Slider(
label="True guidance scale",
minimum=1.0,
maximum=10.0,
step=0.1,
value=1.0
)
num_inference_steps = gr.Slider(
label="Number of inference steps",
minimum=1,
maximum=40,
step=1,
value=4,
)
height = gr.Slider(
label="Height",
minimum=256,
maximum=2048,
step=8,
value=None, value=None,
interactive=True
)
refresh_btn = gr.Button("🔄", size="sm", scale=0)
with gr.Accordion("Settings", open=False):
h_slider = gr.Slider(512, 2048, 1024, step=64, label="Height")
w_slider = gr.Slider(512, 2048, 1024, step=64, label="Width")
steps_slider = gr.Slider(1, 50, 9, step=1, label="Steps")
seed_num = gr.Number(42, label="Seed")
rand_seed = gr.Checkbox(True, label="Randomize Seed")
gen_btn = gr.Button("Generate", variant="primary")
with gr.Column():
out_img = gr.Image(label="Result")
out_seed = gr.Number(label="Seed Used")
# TAB 2: TRAINING
with gr.Tab("🏋️ Train LoRA"):
gr.Markdown("⚠️ **Note:** Requires paid GPU space for long timeouts.")
with gr.Row():
with gr.Column():
train_files = gr.Files(label="Images", file_types=["image"])
train_name = gr.Textbox(label="Project Name", value="my_lora")
train_trigger = gr.Textbox(label="Trigger Word", value="ohwx")
# hidden state for dataset path
dataset_path_state = gr.State()
upload_btn = gr.Button("1. Process Dataset")
upload_status = gr.Textbox(label="Dataset Status")
gr.Markdown("---")
train_steps = gr.Slider(100, 2000, 500, step=100, label="Steps")
train_lr = gr.Slider(1e-5, 1e-3, 1e-4, step=1e-5, label="Learning Rate")
train_rank = gr.Slider(4, 128, 16, step=4, label="Rank")
start_train_btn = gr.Button("2. Start Training", variant="stop")
with gr.Column():
train_log = gr.Textbox(label="Training Log", lines=10)
lora_file_download = gr.File(label="Download LoRA")
# WIRING
# Refresh LoRA list
refresh_btn.click(update_lora_list, outputs=lora_selector)
# Upload
upload_btn.click(
upload_and_prepare_dataset,
[train_files, train_name, train_trigger],
[upload_status, dataset_path_state, train_name]
) )
# Train width = gr.Slider(
def on_train_complete(status, file_path): label="Width",
# Update available loras list immediately after training minimum=256,
new_choices = [("None", None)] + [(k, v) for k, v in AVAILABLE_LORAS.items()] maximum=2048,
return status, file_path, gr.Dropdown(choices=new_choices) step=8,
value=None,
start_train_btn.click(
train_lora,
[dataset_path_state, train_name, train_trigger, train_steps, train_lr, train_rank, h_slider], # reusing h_slider for res
[train_log, lora_file_download]
).then(
update_lora_list,
outputs=[lora_selector]
) )
# Generate
gen_btn.click( rewrite_prompt = gr.Checkbox(label="Rewrite prompt", value=True)
generate_image,
[prompt_input, h_slider, w_slider, steps_slider, seed_num, rand_seed, lora_selector, train_lr], # train_lr dummy # gr.Examples(examples=examples, inputs=[prompt], outputs=[result, seed], fn=infer, cache_examples=False)
[out_img, out_seed]
gr.on(
triggers=[run_button.click, prompt.submit],
fn=infer,
inputs=[
input_images,
prompt,
seed,
randomize_seed,
true_guidance_scale,
num_inference_steps,
height,
width,
rewrite_prompt,
],
outputs=[result, seed, use_output_btn], # Added use_output_btn to outputs
)
# Add the new event handler for the "Use Output as Input" button
use_output_btn.click(
fn=use_output_as_input,
inputs=[result],
outputs=[input_images]
) )
if __name__ == "__main__": if __name__ == "__main__":
demo.launch() demo.launch(mcp_server=True)