Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import subprocess | |
| import shutil | |
| import json | |
| import time | |
| from pathlib import Path | |
| import torch | |
| # Setup directories | |
| DATASET_DIR = Path("./datasets") | |
| OUTPUT_DIR = Path("./output") | |
| DATASET_DIR.mkdir(exist_ok=True) | |
| OUTPUT_DIR.mkdir(exist_ok=True) | |
| # Global variable to store dataset path | |
| current_dataset_path = None | |
| def check_gpu(): | |
| """Check if GPU is available""" | |
| if torch.cuda.is_available(): | |
| gpu_name = torch.cuda.get_device_name(0) | |
| return f"✅ GPU Available: {gpu_name}" | |
| return "⚠️ No GPU detected - training will be slow" | |
| def upload_and_prepare_dataset(files, dataset_name, trigger_word): | |
| """Upload images and prepare dataset""" | |
| global current_dataset_path | |
| if not files: | |
| return "❌ Please upload at least one image", None, "" | |
| if not dataset_name: | |
| dataset_name = f"dataset_{int(time.time())}" | |
| # Create dataset directory | |
| dataset_path = DATASET_DIR / dataset_name | |
| dataset_path.mkdir(exist_ok=True, parents=True) | |
| # Save images | |
| image_count = 0 | |
| for file in files: | |
| if file.name.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.bmp')): | |
| filename = Path(file.name).name | |
| destination = dataset_path / filename | |
| shutil.copy(file.name, destination) | |
| # Create simple caption file | |
| caption_file = destination.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. Upload PNG, JPG, JPEG, or WEBP files.", None, "" | |
| current_dataset_path = str(dataset_path) | |
| status = f"✅ Successfully uploaded {image_count} images\n" | |
| status += f"📁 Dataset: {dataset_name}\n" | |
| if trigger_word: | |
| status += f"🏷️ Trigger word: '{trigger_word}'\n" | |
| status += f"💾 Location: {current_dataset_path}" | |
| return status, current_dataset_path, f"Dataset ready: {dataset_name}" | |
| def train_lora( | |
| dataset_path, | |
| project_name, | |
| trigger_word, | |
| steps, | |
| learning_rate, | |
| lora_rank, | |
| resolution, | |
| progress=gr.Progress() | |
| ): | |
| """Train LoRA model""" | |
| if not dataset_path or not os.path.exists(dataset_path): | |
| return "❌ Please upload a dataset first!", 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) | |
| # Create training config | |
| 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": max(100, int(steps / 4)), | |
| "max_step_saves_to_keep": 3, | |
| }, | |
| "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, | |
| }, | |
| "sample": { | |
| "sampler": "flowmatch", | |
| "sample_every": max(100, int(steps / 4)), | |
| "width": int(resolution), | |
| "height": int(resolution), | |
| "prompts": [ | |
| f"{trigger_word} high quality photo" if trigger_word else "high quality photo", | |
| f"{trigger_word} beautiful scene" if trigger_word else "beautiful scene", | |
| ], | |
| "neg": "", | |
| "seed": 42, | |
| "guidance_scale": 0.0, | |
| "sample_steps": 9, | |
| }, | |
| }] | |
| } | |
| } | |
| # Save config | |
| config_path = output_path / "config.json" | |
| with open(config_path, 'w') as f: | |
| json.dump(config, f, indent=2) | |
| progress(0.1, desc="Installing AI Toolkit...") | |
| # Install AI Toolkit if not exists | |
| if not Path("./ai-toolkit").exists(): | |
| try: | |
| subprocess.run( | |
| ["git", "clone", "https://github.com/ostris/ai-toolkit.git"], | |
| check=True, | |
| capture_output=True | |
| ) | |
| os.chdir("ai-toolkit") | |
| subprocess.run( | |
| ["git", "submodule", "update", "--init", "--recursive"], | |
| check=True, | |
| capture_output=True | |
| ) | |
| subprocess.run( | |
| ["pip", "install", "-q", "-r", "requirements.txt"], | |
| check=True | |
| ) | |
| os.chdir("..") | |
| except Exception as e: | |
| return f"❌ Failed to install AI Toolkit: {str(e)}", None | |
| progress(0.3, desc="Starting training...") | |
| # Run training | |
| try: | |
| result = subprocess.run( | |
| ["python", "ai-toolkit/run.py", str(config_path)], | |
| capture_output=True, | |
| text=True, | |
| timeout=3600 # 1 hour timeout | |
| ) | |
| if result.returncode != 0: | |
| return f"❌ Training failed:\n{result.stderr}", None | |
| progress(0.9, desc="Training complete! Finding LoRA file...") | |
| # Find the trained LoRA file | |
| lora_files = list(output_path.glob("*.safetensors")) | |
| if lora_files: | |
| lora_file = lora_files[-1] # Get the latest one | |
| success_msg = f"✅ Training Complete!\n\n" | |
| success_msg += f"📦 LoRA saved: {lora_file.name}\n" | |
| success_msg += f"💾 Size: {lora_file.stat().st_size / (1024*1024):.2f} MB\n" | |
| success_msg += f"🏷️ Use trigger word: '{trigger_word}' in your prompts" | |
| return success_msg, str(lora_file) | |
| else: | |
| return "⚠️ Training completed but no LoRA file found", None | |
| except subprocess.TimeoutExpired: | |
| return "❌ Training timeout (> 1 hour). Try reducing steps.", None | |
| except Exception as e: | |
| return f"❌ Training error: {str(e)}", None | |
| # Gradio Interface | |
| with gr.Blocks(title="Z-Image LoRA Trainer", theme=gr.themes.Soft()) as demo: | |
| gr.Markdown(""" | |
| # 🎨 Z-Image LoRA Trainer | |
| Train custom LoRA models for Z-Image-Base (6B parameter model) | |
| **Quick Start:** | |
| 1. Upload 10-50 images of your subject | |
| 2. Enter a trigger word (e.g., "mycharacter", "mystyle") | |
| 3. Click Train | |
| 4. Download your LoRA when complete | |
| ⚠️ **Note:** Training takes 10-30 minutes depending on steps. Don't close this tab! | |
| """) | |
| # GPU Status | |
| gpu_status = gr.Textbox(label="GPU Status", value=check_gpu(), interactive=False) | |
| with gr.Tab("📤 Upload Dataset"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| file_input = gr.Files( | |
| label="Upload Images (10-50 recommended)", | |
| file_types=["image"], | |
| file_count="multiple" | |
| ) | |
| dataset_name_input = gr.Textbox( | |
| label="Dataset Name", | |
| placeholder="my_dataset", | |
| value="my_dataset" | |
| ) | |
| trigger_word_input = gr.Textbox( | |
| label="Trigger Word (optional but recommended)", | |
| placeholder="e.g., mycharacter, mystyle", | |
| info="A unique word to activate your LoRA" | |
| ) | |
| upload_btn = gr.Button("📤 Upload Dataset", variant="primary", size="lg") | |
| with gr.Column(): | |
| upload_status = gr.Textbox(label="Upload Status", lines=8) | |
| dataset_path_state = gr.Textbox(label="Dataset Path", visible=False) | |
| dataset_ready = gr.Textbox(label="Ready to Train", interactive=False) | |
| with gr.Tab("🚀 Train LoRA"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| project_name_input = gr.Textbox( | |
| label="Project Name", | |
| placeholder="my_lora", | |
| value="my_lora" | |
| ) | |
| gr.Markdown("### Training Settings") | |
| steps_input = gr.Slider( | |
| label="Training Steps", | |
| minimum=100, | |
| maximum=3000, | |
| value=1000, | |
| step=100, | |
| info="More steps = better quality but slower. Start with 1000." | |
| ) | |
| learning_rate_input = gr.Slider( | |
| label="Learning Rate", | |
| minimum=0.00001, | |
| maximum=0.001, | |
| value=0.0001, | |
| step=0.00001, | |
| info="Default 0.0001 works well for most cases" | |
| ) | |
| lora_rank_input = gr.Slider( | |
| label="LoRA Rank", | |
| minimum=4, | |
| maximum=128, | |
| value=16, | |
| step=4, | |
| info="Higher = more detail but larger file. 16 is balanced." | |
| ) | |
| resolution_input = gr.Radio( | |
| label="Resolution", | |
| choices=[512, 768, 1024], | |
| value=1024, | |
| info="Z-Image native resolution is 1024x1024" | |
| ) | |
| train_btn = gr.Button("🚀 Start Training", variant="primary", size="lg") | |
| with gr.Column(): | |
| training_status = gr.Textbox(label="Training Status", lines=15) | |
| lora_output = gr.File(label="Download Trained LoRA") | |
| with gr.Tab("ℹ️ Help"): | |
| gr.Markdown(""" | |
| ## 📚 How to Use | |
| ### Step 1: Prepare Your Images | |
| - **10-50 images** of your subject (more is better for complex subjects) | |
| - **Consistent subject** across images | |
| - **Good variety** in poses, angles, lighting | |
| - **High quality** photos (clear, well-lit) | |
| ### Step 2: Upload Dataset | |
| - Choose a descriptive **dataset name** | |
| - Add a **trigger word** (e.g., "sks person", "mystyle") | |
| - Upload your images | |
| ### Step 3: Configure Training | |
| - **Project name**: Name for your LoRA | |
| - **Steps**: | |
| - 500-1000 for simple subjects | |
| - 1000-2000 for complex subjects/styles | |
| - **Learning rate**: Keep default (0.0001) | |
| - **LoRA Rank**: 16 is good for most cases | |
| ### Step 4: Train | |
| - Click "Start Training" | |
| - Wait 10-30 minutes (don't close tab) | |
| - Download your LoRA when complete | |
| ### Step 5: Use Your LoRA | |
| - Load in ComfyUI, Automatic1111, or other Z-Image tools | |
| - Use your trigger word in prompts | |
| - Example: "a photo of [trigger_word] in a forest" | |
| ## 🎯 Tips for Best Results | |
| - **Good dataset** = good results | |
| - **Consistent subject** across images | |
| - **Unique trigger word** (not common words) | |
| - **Start with 1000 steps**, adjust if needed | |
| - **Don't overtrain** (if quality decreases, reduce steps) | |
| ## ⚠️ Troubleshooting | |
| **Training fails with OOM error:** | |
| - Reduce resolution to 768 or 512 | |
| - Use fewer steps | |
| - Upload fewer images | |
| **LoRA doesn't look like subject:** | |
| - Upload more images (20-30+) | |
| - Increase steps to 1500-2000 | |
| - Ensure images are consistent | |
| **LoRA is too strong/weak:** | |
| - Adjust LoRA weight in your inference tool (0.5-1.5) | |
| ## 📖 Resources | |
| - **Z-Image Model**: [Tongyi-MAI/Z-Image-Base](https://huggingface.co/Tongyi-MAI/Z-Image-Base) | |
| - **AI Toolkit**: [github.com/ostris/ai-toolkit](https://github.com/ostris/ai-toolkit) | |
| - **Training Adapter**: [ostris/zimage_turbo_training_adapter](https://huggingface.co/ostris/zimage_turbo_training_adapter) | |
| """) | |
| # Event handlers | |
| upload_btn.click( | |
| fn=upload_and_prepare_dataset, | |
| inputs=[file_input, dataset_name_input, trigger_word_input], | |
| outputs=[upload_status, dataset_path_state, dataset_ready] | |
| ) | |
| train_btn.click( | |
| fn=train_lora, | |
| inputs=[ | |
| dataset_path_state, | |
| project_name_input, | |
| trigger_word_input, | |
| steps_input, | |
| learning_rate_input, | |
| lora_rank_input, | |
| resolution_input | |
| ], | |
| outputs=[training_status, lora_output] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |