update

.gitignore

@@ -0,0 +1,55 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage
+htmlcov/
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# Local secrets and machine-specific configuration
+.env
+.env.*
+!.env.example
+*.local
+
+# IDE and operating-system files
+.idea/
+.vscode/
+.DS_Store
+Thumbs.db
+
+# Hugging Face, Gradio, and experiment caches
+.cache/
+huggingface/
+gradio_cached_examples/
+wandb/
+runs/
+logs/
+
+# Helpers downloaded lazily by app.py
+/train_dreambooth_lora_sdxl_advanced.py
+/script.py
+
+# Temporary dataset and SpaceRunner directories (UUID names)
+/????????-????-????-????-????????????/
+
+# Generated training artifacts
+outputs/
+output/
+checkpoint-*/
+*.safetensors
+*.ckpt
+*.pt
+*.pth
+
+# Build and packaging artifacts
+build/
+dist/
+*.egg-info/

README.md

@@ -14,7 +14,100 @@ hf_oauth_scopes:
 - read-repos
 - write-repos
 - manage-repos
+models:
+- microsoft/Mage-VL
 short_description: Train LoRAs with Ease
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+# LoRA Ease: SDXL training on Hugging Face
+
+This fork trains LoRA adapters for `stabilityai/stable-diffusion-xl-base-1.0`.
+The maintained Hugging Face Diffusers v0.39.0 advanced SDXL trainer is downloaded
+from a pinned commit the first time training starts. Captions are generated by a
+dedicated Mage-VL Space, while the face-prior dataset is loaded only when needed.
+
+## Deploy as a Hugging Face Space
+
+Create or duplicate a Gradio Space and push these files to it. The UI itself can
+remain on CPU Basic: when a user starts training, AutoTrain SpaceRunner creates a
+temporary L40S training Space under that user's account and publishes the finished
+LoRA to their model repository.
+
+Captioning defaults to `lienthealien/mage-vl-demo`, which uses the Apache-2.0
+`microsoft/Mage-VL` model. If that Space is private, add
+`MAGE_VL_TOKEN` as a secret in this Space's settings. It only needs permission to
+call the Mage-VL Space. `MAGE_VL_SPACE_ID` and `MAGE_VL_API_NAME` can be set as
+Space variables if you use a different duplicate or endpoint. Uploaded images are
+sent to that configured Space for captioning, so use a Space you trust.
+
+The user must:
+
+1. Have a Hugging Face account with a payment method configured.
+2. Accept the SDXL base-model license on their Hugging Face account.
+3. Provide a Hugging Face **write** token when prompted. It is passed to the
+   temporary trainer and is not written into this repository.
+
+Training currently targets the `l40sx1` hardware flavor (48 GB VRAM). The UI's
+cost display is an estimate; Hugging Face billing and hardware pricing are the
+source of truth.
+
+## Optional local development
+
+Local mode remains useful for testing the UI or training on your own NVIDIA GPU.
+Use Python 3.10 or 3.11, create a virtual environment, install a CUDA-enabled
+PyTorch build appropriate for your driver, and then install the project:
+
+```powershell
+py -3.11 -m venv .venv
+.venv\Scripts\Activate.ps1
+# Install CUDA PyTorch first: https://pytorch.org/get-started/locally/
+pip install -r requirements_local.txt
+python app.py
+```
+
+SDXL training is memory-intensive. The UI defaults to batch size 1, gradient
+accumulation 4, gradient checkpointing, latent caching, fp16, and 1024px. For local
+out-of-memory errors, keep batch size 1 and try 768px; 1024px is preferred when the
+GPU permits.
+
+The base model is gated. Before the first run, accept its license on Hugging Face
+and authenticate locally with `hf auth login`.
+
+## Training a Daggerfall-like face style
+
+If the images show many different people sharing the same visual language, this is
+a **style LoRA**, not a face/identity LoRA:
+
+1. Prepare roughly 20-60 varied, clean portraits. Remove menus, text, cursors,
+   duplicates, and unrelated UI. Do not mix modern fan art with original game art
+   unless you deliberately want both aesthetics.
+2. Crop loosely around each head and shoulders, keeping several face angles,
+   ancestries, genders, expressions, backgrounds, and lighting conditions. Do not
+   enlarge tiny source images with heavy smoothing; crisp nearest-neighbor or mild
+   restoration usually preserves the pixel-art language better.
+3. In the UI choose **style** and keep the concept sentence
+   `in the style of TOK`. `TOK` is the trainer's required token abstraction; do
+   not replace it with `daggerfall` or another common word.
+4. Caption what should remain controllable, for example:
+   `portrait of an elderly elf man, stern expression, dark background, in the style of TOK`.
+   Keep the trigger in every caption, but vary subject traits accurately.
+5. Start with rank 32 and 800-1500 steps. Save/test intermediate versions if you
+   raise the step count; overtraining shows up as every output copying the training
+   backgrounds, poses, or facial proportions.
+
+After uploading the images, click **Caption all images with Mage-VL**. The portrait
+instruction asks for controllable character traits while deliberately excluding
+the shared art style; the concept sentence is appended automatically. Review the
+generated text boxes and correct names, species, clothing, or expressions that the
+captioner guessed incorrectly. **Export captions** downloads the reviewed entries
+as ImageFolder-compatible `metadata.jsonl` for backup or offline editing.
+
+Use **character** only if all images are the same fictional character. Use **face**
+only for one consistent real or fictional identity; that preset enables prior
+preservation and is not appropriate for learning a whole game's portrait style.
+
+The output folder contains `pytorch_lora_weights.safetensors`, which can be loaded
+as an SDXL LoRA by Diffusers, ComfyUI, or compatible Stable Diffusion frontends.
+Only train and distribute assets you have the rights or permission to use.

app.py

@@ -1,42 +1,79 @@
 import gradio as gr
-from PIL import Image
-import requests
 import subprocess
 import os
 is_spaces = True if os.environ.get('SPACE_ID') else False
-if(is_spaces):
-    import spaces
-from transformers import Blip2Processor, Blip2ForConditionalGeneration
+from gradio_client import Client, handle_file
 from huggingface_hub import snapshot_download, HfApi
-import torch
 import uuid
 import shutil
 import json
 import random
+import hashlib
 from slugify import slugify
 import argparse 
 import importlib
 import sys
 from pathlib import Path
+import tempfile
 import zipfile
+from urllib.request import urlretrieve
 
 MAX_IMAGES = 150
 
-training_script_url = "https://raw.githubusercontent.com/huggingface/diffusers/ba28006f8b2a0f7ec3b6784695790422b4f80a97/examples/advanced_diffusion_training/train_dreambooth_lora_sdxl_advanced.py"
-subprocess.run(['wget', '-N', training_script_url])
+DIFFUSERS_VERSION = "v0.39.0"
+DIFFUSERS_COMMIT = "a3608b512ed7248499a44c61d954965ed9bdae4d"
+TRAINING_SCRIPT = Path("train_dreambooth_lora_sdxl_advanced.py")
+training_script_url = f"https://raw.githubusercontent.com/huggingface/diffusers/{DIFFUSERS_COMMIT}/examples/advanced_diffusion_training/{TRAINING_SCRIPT.name}"
 orchestrator_script_url = "https://huggingface.co/datasets/multimodalart/lora-ease-helper/raw/main/script.py"
-subprocess.run(['wget', '-N', orchestrator_script_url])
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+MAGE_VL_SPACE_ID = os.environ.get("MAGE_VL_SPACE_ID", "lienthealien/mage-vl-demo")
+MAGE_VL_API_NAME = os.environ.get("MAGE_VL_API_NAME", "/ask_image")
+mage_vl_client = None
+caption_cache = {}
 
-FACES_DATASET_PATH = snapshot_download(repo_id="multimodalart/faces-prior-preservation", repo_type="dataset")
-#Delete .gitattributes to process things properly
-Path(FACES_DATASET_PATH, '.gitattributes').unlink(missing_ok=True)
 
-processor = Blip2Processor.from_pretrained("Salesforce/blip2-opt-2.7b")
-model = Blip2ForConditionalGeneration.from_pretrained(
-    "Salesforce/blip2-opt-2.7b", device_map={"": 0}, torch_dtype=torch.float16
-)
+def ensure_file(url, destination):
+    """Download a pinned helper only when it is actually needed."""
+    destination = Path(destination)
+    if not destination.exists():
+        print(f"Downloading {destination.name} from {url}")
+        urlretrieve(url, destination)
+    return destination
+
+
+def get_face_prior_dataset():
+    dataset_path = snapshot_download(
+        repo_id="multimodalart/faces-prior-preservation", repo_type="dataset"
+    )
+    Path(dataset_path, ".gitattributes").unlink(missing_ok=True)
+    return dataset_path
+
+
+def get_captioner():
+    """Connect to the dedicated Mage-VL Space without loading a VLM here."""
+    global mage_vl_client
+    if mage_vl_client is None:
+        token = os.environ.get("MAGE_VL_TOKEN") or os.environ.get("HF_TOKEN")
+        mage_vl_client = Client(MAGE_VL_SPACE_ID, hf_token=token, verbose=False)
+    return mage_vl_client
+
+
+def caption_with_mage(client, image, instruction):
+    with open(image, "rb") as image_file:
+        cache_key = hashlib.sha256(image_file.read() + instruction.encode("utf-8")).hexdigest()
+    if cache_key in caption_cache:
+        return caption_cache[cache_key]
+
+    result = client.predict(
+        handle_file(image),
+        instruction,
+        160,
+        api_name=MAGE_VL_API_NAME,
+    )
+    caption = result[0] if isinstance(result, (list, tuple)) else result
+    caption = str(caption).strip().rstrip(" .,")
+    caption_cache[cache_key] = caption
+    return caption
 
 training_option_settings = {
     "face": {
@@ -46,7 +83,7 @@ training_option_settings = {
         "class_prompt": "a photo of a person",
         "train_steps_multiplier": 75,
         "file_count": 150,
-        "dataset_path": FACES_DATASET_PATH
+        "dataset_path": None
     },
     "style": {
         "rank": 32,
@@ -156,7 +193,7 @@ def change_defaults(option, images):
 
     random_files = []
     if settings["with_prior_preservation"]:
-        directory = settings["dataset_path"]
+        directory = settings["dataset_path"] or get_face_prior_dataset()
         file_count = settings["file_count"]
         files = [os.path.join(directory, file) for file in os.listdir(directory) if os.path.isfile(os.path.join(directory, file))]
         random_files = random.sample(files, min(len(files), file_count))
@@ -325,21 +362,26 @@ def start_training(
     spacerunner_args = ';'.join(commands)
     if not os.path.exists(spacerunner_folder):
         os.makedirs(spacerunner_folder)
-    shutil.copy("train_dreambooth_lora_sdxl_advanced.py", f"{spacerunner_folder}/trainer.py")
+    ensure_file(training_script_url, TRAINING_SCRIPT)
+    ensure_file(orchestrator_script_url, "script.py")
+    shutil.copy(TRAINING_SCRIPT, f"{spacerunner_folder}/trainer.py")
     shutil.copy("script.py", f"{spacerunner_folder}/script.py")
     shutil.copytree(dataset_folder, f"{spacerunner_folder}/{dataset_folder}")
-    requirements='''peft==0.7.1
--huggingface_hub
+    requirements='''diffusers==0.39.0
+peft>=0.11.1
 torch
-git+https://github.com/huggingface/diffusers@ba28006f8b2a0f7ec3b6784695790422b4f80a97
-transformers==4.36.2
-accelerate==0.25.0
-safetensors==0.4.1
+torchvision
+transformers>=4.41.2
+accelerate>=0.31.0
+safetensors>=0.4.3
 prodigyopt==1.0
-hf-transfer==0.1.4
-huggingface_hub==0.20.3
-git+https://github.com/huggingface/datasets.git@3f149204a2a5948287adcade5e90707aa5207a92
-numpy==1.26.4'''
+huggingface_hub>=0.34.0
+hf-transfer>=0.1.4
+datasets>=2.20.0
+ftfy
+tensorboard
+Jinja2
+sentencepiece'''
     file_path = f'{spacerunner_folder}/requirements.txt'
     with open(file_path, 'w') as file:
         file.write(requirements)
@@ -495,6 +537,8 @@ def start_training_og(
         commands.append(f"--prodigy_beta3={prodigy_beta3}")
     if use_adam_weight_decay_text_encoder:
         commands.append(f"--adam_weight_decay_text_encoder={adam_weight_decay_text_encoder}")
+    ensure_file(training_script_url, TRAINING_SCRIPT)
+    importlib.invalidate_caches()
     from train_dreambooth_lora_sdxl_advanced import main as train_main, parse_args as parse_train_args
     args = parse_train_args(commands)
     
@@ -503,27 +547,52 @@ def start_training_og(
     return f"Your model has finished training and has been saved to the `{slugged_lora_name}` folder"
 
 def run_captioning(*inputs):
-    model.to("cuda")
     images = inputs[0]
-    training_option = inputs[-1]
+    training_option = inputs[-2]
+    caption_instruction = inputs[-1].strip()
     final_captions = [""] * MAX_IMAGES
+    try:
+        caption_client = get_captioner()
+    except Exception as exc:
+        raise gr.Error(
+            f"Could not connect to Mage-VL Space '{MAGE_VL_SPACE_ID}'. "
+            "If it is private, add MAGE_VL_TOKEN as a Hugging Face Space secret. "
+            f"Details: {exc}"
+        ) from exc
+
     for index, image in enumerate(images):
-        original_caption = inputs[index + 1]
-        pil_image = Image.open(image)  
-        blip_inputs = processor(images=pil_image, return_tensors="pt").to(device, torch.float16)
-        generated_ids = model.generate(**blip_inputs)
-        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
+        concept_caption = inputs[index + 1].strip()
+        try:
+            generated_text = caption_with_mage(caption_client, image, caption_instruction)
+        except Exception as exc:
+            raise gr.Error(
+                f"Mage-VL failed while captioning image {index + 1}: {exc}"
+            ) from exc
         if training_option == "style":
-            final_caption = generated_text + " " + original_caption
+            final_caption = f"{generated_text}, {concept_caption}"
         else:
-            final_caption = original_caption + " " + generated_text
+            final_caption = f"{concept_caption}, {generated_text}"
         final_captions[index] = final_caption
         yield final_captions
 
-if is_spaces:
-    run_captioning = spaces.GPU()(run_captioning)
 
-        
+def export_captions(images, *captions):
+    if not images:
+        raise gr.Error("Upload images before exporting captions.")
+
+    with tempfile.NamedTemporaryFile(
+        mode="w", suffix=".jsonl", prefix="lora-captions-", encoding="utf-8", delete=False
+    ) as caption_file:
+        for image, caption in zip(images, captions):
+            caption_file.write(
+                json.dumps(
+                    {"file_name": os.path.basename(image), "prompt": (caption or "").strip()},
+                    ensure_ascii=False,
+                )
+                + "\n"
+            )
+    return caption_file.name
+
 def check_token(token):
     try:
         api = HfApi(token=token)
@@ -578,7 +647,9 @@ with gr.Blocks(css=css, theme=theme) as demo:
         with gr.Column(elem_classes=["main_logged"]) as main_ui:
             lora_name = gr.Textbox(label="The name of your LoRA", info="This has to be a unique name", placeholder="e.g.: Persian Miniature Painting style, Cat Toy")
             training_option = gr.Radio(
-                label="What are you training?", choices=["object", "style", "character", "face", "custom"]
+                label="What are you training?",
+                info="For a shared game look across many different faces, choose style. Choose character only when every image depicts the same character.",
+                choices=["object", "style", "character", "face", "custom"],
             )
             concept_sentence = gr.Textbox(
                 label="Concept sentence",
@@ -604,7 +675,19 @@ with gr.Blocks(css=css, theme=theme) as demo:
 To improve the quality of your outputs, you can add a custom caption for each image, describing exactly what is taking place in each of them. Including TOK is mandatory. You can leave things as is if you don't want to include captioning.
 """
                             )
-                            do_captioning = gr.Button("Add AI captions with BLIP-2")
+                            caption_instruction = gr.Textbox(
+                                label="AI caption instruction",
+                                value="Write one concise SDXL training caption for this portrait. Describe only visible, controllable content: character species or ancestry, approximate age, gender presentation, expression, hair, clothing, pose, lighting, and background. Do not identify the game, artist, rendering technique, image quality, or art style. Do not speculate about details that are not visible. Return only a natural comma-separated caption with no introduction.",
+                                info="The LoRA concept sentence is added automatically to every generated caption.",
+                                lines=3,
+                            )
+                            with gr.Row():
+                                do_captioning = gr.Button("Caption all images with Mage-VL", variant="primary")
+                                export_caption_button = gr.Button("Export captions")
+                            exported_captions = gr.File(
+                                label="Caption metadata (.jsonl)",
+                                interactive=False,
+                            )
                             output_components = [captioning_area]
                             caption_list = []
                             for i in range(1, MAX_IMAGES + 1):
@@ -653,7 +736,7 @@ To improve the quality of your outputs, you can add a custom caption for each im
                         mixed_precision = gr.Dropdown(
                             label="Mixed Precision",
                             choices=["no", "fp16", "bf16"],
-                            value="bf16",
+                            value="fp16",
                         )
                         learning_rate = gr.Number(
                             label="UNet Learning rate",
@@ -753,12 +836,12 @@ To improve the quality of your outputs, you can add a custom caption for each im
                             gradient_accumulation_steps = gr.Number(
                                 info="If you change this setting, the pricing calculation will be wrong",
                                 label="gradient_accumulation_steps", 
-                                value=1
+                                value=4
                             )
                             train_batch_size = gr.Number(
                                 info="If you change this setting, the pricing calculation will be wrong",
                                 label="Train batch size",
-                                value=2
+                                value=1
                             )
                             num_train_epochs = gr.Number(
                                 info="If you change this setting, the pricing calculation will be wrong",
@@ -1070,9 +1153,17 @@ If you prefer command line, you can run our [training script]({training_script_u
     )
 
     do_captioning.click(
-        fn=run_captioning, inputs=[images] + caption_list + [training_option], outputs=caption_list
+        fn=run_captioning,
+        inputs=[images] + caption_list + [training_option, caption_instruction],
+        outputs=caption_list,
+    )
+    export_caption_button.click(
+        fn=export_captions,
+        inputs=[images] + caption_list,
+        outputs=exported_captions,
+        queue=False,
     )
     #demo.load(fn=swap_opacity, outputs=[main_ui], queue=False, concurrency_limit=50)
 if __name__ == "__main__":
     demo.queue()
-    demo.launch(share=True)
+    demo.launch(share=True)

requirements.txt

@@ -1,9 +1,6 @@
 autotrain-advanced
 numpy==1.26.4
-transformers
-huggingface-hub
+huggingface-hub==0.27.0
 python-slugify
-uuid
-torch
-torchvision
-spaces
+gradio>=4.37,<6
+gradio_client>=1.0

requirements_local.txt

@@ -1,14 +1,16 @@
 torch
-python-slugify
-uuid
-torch
 torchvision
-peft==0.7.1
-huggingface-hub==0.23.4
-diffusers==0.29.2
-transformers==4.42.3
-accelerate==0.31.0
-safetensors==0.4.3
+python-slugify
+gradio>=4.37,<6
+diffusers==0.39.0
+peft>=0.11.1
+huggingface-hub>=0.34.0
+transformers>=4.41.2
+accelerate>=0.31.0
+safetensors>=0.4.3
 prodigyopt==1.0
-hf-transfer==0.1.4
-datasets==2.20.0
+datasets>=2.20.0
+ftfy
+tensorboard
+Jinja2
+sentencepiece