diff --git a/app.py b/app.py new file mode 100644 index 0000000..5218940 --- /dev/null +++ b/app.py @@ -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() \ No newline at end of file