Create app.py
This commit is contained in:
sigma
340
app.py
Normal file
340
app.py
Normal file
@@ -0,0 +1,340 @@
|
||||
import gradio as gr
|
||||
import os
|
||||
import subprocess
|
||||
import shutil
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
import torch
|
||||
import spaces
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
# ==========================================
|
||||
# 1. SETUP & GLOBAL VARS
|
||||
# ==========================================
|
||||
|
||||
DATASET_DIR = Path("./datasets")
|
||||
OUTPUT_DIR = Path("./output")
|
||||
DATASET_DIR.mkdir(exist_ok=True)
|
||||
OUTPUT_DIR.mkdir(exist_ok=True)
|
||||
|
||||
# global tracking for loras
|
||||
# key: friendly name, value: path
|
||||
AVAILABLE_LORAS = {}
|
||||
|
||||
print("loading z-image-turbo pipeline...")
|
||||
pipe = DiffusionPipeline.from_pretrained(
|
||||
"Tongyi-MAI/Z-Image-Turbo",
|
||||
torch_dtype=torch.bfloat16,
|
||||
low_cpu_mem_usage=False,
|
||||
)
|
||||
pipe.to("cuda")
|
||||
print("pipeline loaded!")
|
||||
|
||||
# ==========================================
|
||||
# 2. TRAINING LOGIC
|
||||
# ==========================================
|
||||
|
||||
def check_gpu():
|
||||
if torch.cuda.is_available():
|
||||
return f"✅ gpu available: {torch.cuda.get_device_name(0)}"
|
||||
return "⚠️ no gpu detected"
|
||||
|
||||
def upload_and_prepare_dataset(files, dataset_name, trigger_word):
|
||||
if not files:
|
||||
return "❌ upload images first", None, ""
|
||||
|
||||
if not dataset_name:
|
||||
dataset_name = f"dataset_{int(time.time())}"
|
||||
|
||||
dataset_path = DATASET_DIR / dataset_name
|
||||
dataset_path.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
image_count = 0
|
||||
for file in files:
|
||||
if file.name.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.bmp')):
|
||||
filename = Path(file.name).name
|
||||
dest = dataset_path / filename
|
||||
shutil.copy(file.name, dest)
|
||||
|
||||
caption_file = dest.with_suffix('.txt')
|
||||
caption_text = trigger_word if trigger_word else "a photo"
|
||||
with open(caption_file, 'w') as f:
|
||||
f.write(caption_text)
|
||||
|
||||
image_count += 1
|
||||
|
||||
if image_count == 0:
|
||||
return "❌ no valid images found", None, ""
|
||||
|
||||
return f"✅ ready: {image_count} images in {dataset_name}", str(dataset_path), dataset_name
|
||||
|
||||
# request 1 hour gpu for training
|
||||
@spaces.GPU(duration=3600)
|
||||
def train_lora(
|
||||
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:
|
||||
project_name = f"lora_{int(time.time())}"
|
||||
|
||||
output_path = OUTPUT_DIR / project_name
|
||||
output_path.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
# config generation
|
||||
config = {
|
||||
"job": "extension",
|
||||
"config": {
|
||||
"name": project_name,
|
||||
"process": [{
|
||||
"type": "sd_trainer",
|
||||
"training_folder": str(output_path),
|
||||
"device": "cuda:0",
|
||||
"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:
|
||||
subprocess.run(["git", "clone", "https://github.com/ostris/ai-toolkit.git"], check=True)
|
||||
subprocess.run(["pip", "install", "-q", "-r", "ai-toolkit/requirements.txt"], check=True)
|
||||
except Exception as e:
|
||||
return f"❌ setup failed: {e}", None
|
||||
|
||||
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:
|
||||
return f"❌ training crashed:\n{proc.stderr}", None
|
||||
|
||||
# find result
|
||||
lora_files = list(output_path.glob("*.safetensors"))
|
||||
if lora_files:
|
||||
lora_file = lora_files[-1]
|
||||
AVAILABLE_LORAS[project_name] = str(lora_file)
|
||||
|
||||
# update the dropdown choices dynamically
|
||||
choices = [("None", None)] + [(k, v) for k, v in AVAILABLE_LORAS.items()]
|
||||
|
||||
return f"✅ trained: {project_name}", str(lora_file)
|
||||
|
||||
return "⚠️ finished but no safetensors found", None
|
||||
|
||||
except Exception as e:
|
||||
return f"❌ fatal error: {e}", None
|
||||
|
||||
# ==========================================
|
||||
# 3. INFERENCE LOGIC
|
||||
# ==========================================
|
||||
|
||||
@spaces.GPU
|
||||
def generate_image(
|
||||
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):
|
||||
print(f"loading lora: {lora_path}")
|
||||
try:
|
||||
pipe.load_lora_weights(lora_path)
|
||||
# manual scaling not always supported directly without fuse,
|
||||
# but usually applied by default.
|
||||
# for simplicitly we just load it.
|
||||
except Exception as e:
|
||||
print(f"lora load failed: {e}")
|
||||
|
||||
if randomize_seed:
|
||||
seed = torch.randint(0, 2**32 - 1, (1,)).item()
|
||||
|
||||
generator = torch.Generator("cuda").manual_seed(int(seed))
|
||||
|
||||
image = pipe(
|
||||
prompt=prompt,
|
||||
height=int(height),
|
||||
width=int(width),
|
||||
num_inference_steps=int(steps),
|
||||
guidance_scale=0.0,
|
||||
generator=generator,
|
||||
).images[0]
|
||||
|
||||
return image, seed
|
||||
|
||||
def update_lora_list():
|
||||
"""helper to refresh dropdown"""
|
||||
choices = [("None", None)] + [(k, v) for k, v in AVAILABLE_LORAS.items()]
|
||||
return gr.Dropdown(choices=choices)
|
||||
|
||||
# ==========================================
|
||||
# 4. UI CONSTRUCTION
|
||||
# ==========================================
|
||||
|
||||
custom_theme = gr.themes.Soft(primary_hue="yellow", secondary_hue="slate")
|
||||
|
||||
with gr.Blocks(theme=custom_theme, title="Z-Image ZeroGPU Trainer") as demo:
|
||||
|
||||
gr.Markdown("# ⚡ Z-Image-Turbo: Train & Test")
|
||||
|
||||
with gr.Tabs():
|
||||
|
||||
# TAB 1: INFERENCE
|
||||
with gr.Tab("🎨 Generate"):
|
||||
with gr.Row():
|
||||
with gr.Column():
|
||||
prompt_input = gr.Textbox(label="Prompt", lines=3)
|
||||
|
||||
with gr.Row():
|
||||
lora_selector = gr.Dropdown(
|
||||
label="Select LoRA",
|
||||
choices=[("None", 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
|
||||
def on_train_complete(status, file_path):
|
||||
# Update available loras list immediately after training
|
||||
new_choices = [("None", None)] + [(k, v) for k, v in AVAILABLE_LORAS.items()]
|
||||
return status, file_path, gr.Dropdown(choices=new_choices)
|
||||
|
||||
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(
|
||||
generate_image,
|
||||
[prompt_input, h_slider, w_slider, steps_slider, seed_num, rand_seed, lora_selector, train_lr], # train_lr dummy
|
||||
[out_img, out_seed]
|
||||
)
|
||||
|
||||
if __name__ == "__main__":
|
||||
demo.launch()
|
||||
Reference in New Issue
Block a user