app.py

import spaces
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}"
#@spaces.GPU()
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()