Create app.py

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()