fix

README.md

@@ -4,8 +4,9 @@ emoji: 🧞
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 4.44.1
+sdk_version: 5.29.0
 app_file: app.py
+startup_duration_timeout: 1h
 pinned: true
 license: mit
 fullWidth: true
@@ -25,23 +26,21 @@ Check out the configuration reference at https://huggingface.co/docs/hub/spaces-
 
 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.
+from a pinned commit the first time training starts. Captions are generated in this
+Space by the embedded Mage-VL model; 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 the public `microsoft/mage-vl-demo`, which uses the
-Apache-2.0 `microsoft/Mage-VL` model. No repository token is required for reading
-its API. `MAGE_VL_SPACE_ID` and `MAGE_VL_API_NAME` can be set as Space variables
-if you use a different duplicate or endpoint. `MAGE_VL_BASE_URL` can override the
-derived `https://OWNER-SPACE.hf.space` URL. If that override is private, add a token
-with access to it as the `MAGE_VL_TOKEN` Space secret. Uploaded images are sent to
-the configured captioning Space, so use a Space you trust.
+Create or duplicate a Gradio Space, enable **ZeroGPU** hardware, and push these
+files to it. `microsoft/Mage-VL` is loaded directly into this Space and the
+captioning function is decorated with `spaces.GPU`; images never need to be sent
+to a second Space or account. The first startup downloads the model and can take
+several minutes.
+
+When a user starts training, the app directly creates a private dataset repository
+and a temporary private L40S Docker Space under that user's account. This keeps
+AutoTrain's conflicting dependency stack out of the captioning/UI container while
+preserving the same trainer runtime and finished-model upload behavior.
 
 The user must:
 

app.py

@@ -1,12 +1,14 @@
-import gradio as gr
-import subprocess
 import os
-import requests
 is_spaces = True if os.environ.get('SPACE_ID') else False
 if is_spaces:
     import spaces
+import gradio as gr
+import torch
+from PIL import Image
+from transformers import AutoModelForCausalLM, AutoProcessor
 from gradio_client import utils as gradio_client_utils
 from huggingface_hub import snapshot_download, HfApi
+import io
 import uuid
 import shutil
 import json
@@ -29,13 +31,9 @@ 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"
 
-MAGE_VL_SPACE_ID = os.environ.get("MAGE_VL_SPACE_ID", "microsoft/mage-vl-demo")
-MAGE_VL_API_NAME = os.environ.get("MAGE_VL_API_NAME", "/ask_image")
-MAGE_VL_BASE_URL = os.environ.get(
-    "MAGE_VL_BASE_URL",
-    f"https://{MAGE_VL_SPACE_ID.replace('/', '-').lower()}.hf.space",
-).rstrip("/")
-mage_vl_client = None
+MAGE_VL_MODEL_ID = "microsoft/Mage-VL"
+mage_vl_processor = None
+mage_vl_model = None
 caption_cache = {}
 
 
@@ -54,16 +52,6 @@ def _safe_schema_to_python_type(schema, defs):
 gradio_client_utils._json_schema_to_python_type = _safe_schema_to_python_type
 
 
-# A Space configured for ZeroGPU must register at least one GPU function during
-# startup. Captioning itself runs in the separate Mage-VL Space, so this probe is
-# intentionally not placed around run_captioning (which would waste local quota
-# while waiting for remote API requests).
-if is_spaces:
-    @spaces.GPU(duration=1)
-    def zero_gpu_probe():
-        return True
-
-
 def ensure_file(url, destination):
     """Download a pinned helper only when it is actually needed."""
     destination = Path(destination)
@@ -81,90 +69,78 @@ def get_face_prior_dataset():
     return dataset_path
 
 
-class MageVLClient:
-    """Minimal Gradio HTTP client, independent of the local Gradio version."""
-
-    def __init__(self, base_url, api_name, token=None):
-        self.base_url = base_url
-        self.api_name = api_name.strip("/")
-        self.session = requests.Session()
-        if token:
-            self.session.headers.update({"Authorization": f"Bearer {token}"})
-
-        response = self.session.get(f"{self.base_url}/gradio_api/info", timeout=30)
-        response.raise_for_status()
-        endpoints = response.json().get("named_endpoints", {})
-        if f"/{self.api_name}" not in endpoints:
-            raise RuntimeError(
-                f"Mage-VL endpoint '/{self.api_name}' was not found at {self.base_url}."
-            )
-
-    def predict(self, image, instruction, max_new_tokens):
-        with open(image, "rb") as image_file:
-            upload = self.session.post(
-                f"{self.base_url}/gradio_api/upload",
-                files={"files": (os.path.basename(image), image_file)},
-                timeout=60,
-            )
-        upload.raise_for_status()
-        uploaded_path = upload.json()[0]
-
-        payload = {
-            "data": [
-                {
-                    "path": uploaded_path,
-                    "orig_name": os.path.basename(image),
-                    "meta": {"_type": "gradio.FileData"},
-                },
-                instruction,
-                int(max_new_tokens),
-            ]
-        }
-        start = self.session.post(
-            f"{self.base_url}/gradio_api/call/{self.api_name}",
-            json=payload,
-            timeout=30,
-        )
-        start.raise_for_status()
-        event_id = start.json()["event_id"]
-
-        result = self.session.get(
-            f"{self.base_url}/gradio_api/call/{self.api_name}/{event_id}",
-            timeout=180,
+def get_captioner():
+    """Load the embedded Mage-VL model once per Space process."""
+    global mage_vl_processor, mage_vl_model
+    if mage_vl_processor is None or mage_vl_model is None:
+        target_device = "cuda" if is_spaces or torch.cuda.is_available() else "cpu"
+        dtype = torch.bfloat16 if target_device == "cuda" else torch.float32
+        mage_vl_processor = AutoProcessor.from_pretrained(
+            MAGE_VL_MODEL_ID, trust_remote_code=True
         )
-        result.raise_for_status()
-        event = None
-        for line in result.text.splitlines():
-            if line.startswith("event:"):
-                event = line.partition(":")[2].strip()
-            elif line.startswith("data:") and event == "complete":
-                return json.loads(line.partition(":")[2].strip())
-            elif line.startswith("data:") and event == "error":
-                raise RuntimeError(line.partition(":")[2].strip())
-        raise RuntimeError(f"Mage-VL returned no completed result: {result.text[:500]}")
+        mage_vl_model = AutoModelForCausalLM.from_pretrained(
+            MAGE_VL_MODEL_ID,
+            trust_remote_code=True,
+            dtype=dtype,
+            attn_implementation="sdpa",
+        ).to(target_device).eval()
+    return mage_vl_processor, mage_vl_model
 
 
-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")
-        mage_vl_client = MageVLClient(MAGE_VL_BASE_URL, MAGE_VL_API_NAME, token)
-    return mage_vl_client
+def _mage_prompt(processor, instruction):
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image"},
+                {"type": "text", "text": instruction},
+            ],
+        }
+    ]
+    return processor.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
 
 
-def caption_with_mage(client, image, instruction):
+def caption_with_mage(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(image, instruction, 160)
-    caption = result[0] if isinstance(result, (list, tuple)) else result
+    processor, model = get_captioner()
+    with Image.open(image) as source_image:
+        pil_image = source_image.convert("RGB")
+    inputs = processor(
+        text=[_mage_prompt(processor, instruction)],
+        images=[pil_image],
+        return_tensors="pt",
+    )
+    device_inputs = {}
+    for key, value in inputs.items():
+        if not hasattr(value, "to"):
+            continue
+        device_inputs[key] = value.to(model.device)
+        if key == "pixel_values":
+            device_inputs[key] = device_inputs[key].to(model.dtype)
+    with torch.inference_mode():
+        output = model.generate(
+            **device_inputs,
+            max_new_tokens=160,
+            do_sample=False,
+        )
+    new_tokens = output[0, device_inputs["input_ids"].shape[1]:]
+    caption = processor.tokenizer.decode(new_tokens, skip_special_tokens=True)
     caption = str(caption).strip().rstrip(" .,")
     caption_cache[cache_key] = caption
     return caption
 
+
+if is_spaces:
+    # ZeroGPU emulates CUDA during startup, so the weights can be placed once and
+    # transparently backed by a real GPU inside the decorated captioning call.
+    get_captioner()
+
 training_option_settings = {
     "face": {
         "rank": 32,
@@ -311,6 +287,88 @@ def create_dataset(*inputs):
     
     return destination_folder
 
+
+AUTOTRAIN_DOCKERFILE = """FROM huggingface/autotrain-advanced:latest
+
+CMD pip uninstall -y autotrain-advanced && pip install -U autotrain-advanced && autotrain api --port 7860 --host 0.0.0.0
+"""
+
+
+def _commands_to_args(commands):
+    """Convert SpaceRunner's old semicolon arguments into its JSON form."""
+    args = {}
+    for command in commands:
+        key, separator, value = command.partition("=")
+        args[key] = value if separator else ""
+    return args
+
+
+def _launch_training_space(folder, project_name, commands, token):
+    """Create the dataset and L40S trainer Space without installing AutoTrain here."""
+    api = HfApi(token=token)
+    username = api.whoami()["name"]
+    repo_name = f"autotrain-{project_name}"
+    dataset_id = f"{username}/{repo_name}"
+    space_id = f"{username}/{repo_name}"
+
+    api.create_repo(repo_id=dataset_id, repo_type="dataset", private=True)
+    api.upload_folder(folder_path=folder, repo_id=dataset_id, repo_type="dataset")
+
+    params = {
+        "project_name": project_name,
+        "data_path": dataset_id,
+        "username": username,
+        "token": token,
+        "script_path": folder,
+        "env": {},
+        "args": _commands_to_args(commands),
+    }
+    api.create_repo(
+        repo_id=space_id,
+        repo_type="space",
+        space_sdk="docker",
+        space_hardware="l40sx1",
+        private=True,
+    )
+    secrets = {
+        "HF_TOKEN": token,
+        "HF_HUB_ENABLE_HF_TRANSFER": "1",
+        "AUTOTRAIN_USERNAME": username,
+        "PROJECT_NAME": project_name,
+        "TASK_ID": "27",
+        "PARAMS": json.dumps(params),
+        "DATA_PATH": dataset_id,
+    }
+    for key, value in secrets.items():
+        api.add_space_secret(repo_id=space_id, key=key, value=value)
+    api.set_space_sleep_time(repo_id=space_id, sleep_time=604800)
+
+    space_readme = f"""---
+title: {project_name}
+emoji: 🚀
+colorFrom: green
+colorTo: indigo
+sdk: docker
+pinned: false
+tags:
+- autotrain
+duplicated_from: autotrain-projects/autotrain-advanced
+---
+"""
+    api.upload_file(
+        path_or_fileobj=io.BytesIO(space_readme.encode("utf-8")),
+        path_in_repo="README.md",
+        repo_id=space_id,
+        repo_type="space",
+    )
+    api.upload_file(
+        path_or_fileobj=io.BytesIO(AUTOTRAIN_DOCKERFILE.encode("utf-8")),
+        path_in_repo="Dockerfile",
+        repo_id=space_id,
+        repo_type="space",
+    )
+    return username, space_id
+
 def start_training(
     lora_name,
     training_option,
@@ -448,8 +506,6 @@ def start_training(
     if use_adam_weight_decay_text_encoder:
         commands.append(f"adam_weight_decay_text_encoder={adam_weight_decay_text_encoder}")
     print(commands)
-    # Joining the commands with ';' separator for spacerunner format
-    spacerunner_args = ';'.join(commands)
     if not os.path.exists(spacerunner_folder):
         os.makedirs(spacerunner_folder)
     ensure_file(training_script_url, TRAINING_SCRIPT)
@@ -475,18 +531,19 @@ sentencepiece'''
     file_path = f'{spacerunner_folder}/requirements.txt'
     with open(file_path, 'w') as file:
         file.write(requirements)
-    # The subprocess call for autotrain spacerunner
-    api = HfApi(token=token)
-    username = api.whoami()["name"]
-    subprocess_command = ["autotrain", "spacerunner", "--project-name", slugged_lora_name, "--script-path", spacerunner_folder, "--username", username, "--token", token, "--backend", "spaces-l40sx1", "--env",f"HF_TOKEN={token};HF_HUB_ENABLE_HF_TRANSFER=1", "--args", spacerunner_args]
-    outcome = subprocess.run(subprocess_command)
-    if(outcome.returncode == 0):
-        return f"""# Your training has started. 
-## - Training Status: <a href='https://huggingface.co/spaces/{username}/autotrain-{slugged_lora_name}?logs=container'>{username}/autotrain-{slugged_lora_name}</a> <small>(in the logs tab)</small>
+    try:
+        username, space_id = _launch_training_space(
+            spacerunner_folder, slugged_lora_name, commands, token
+        )
+    except Exception as exc:
+        raise gr.Error(
+            "Could not create the private L40S training Space. Make sure the "
+            f"LoRA name is unique and the token can create repositories. Details: {exc}"
+        ) from exc
+
+    return f"""# Your training has started.
+## - Training Status: <a href='https://huggingface.co/spaces/{space_id}?logs=container'>{space_id}</a> <small>(in the logs tab)</small>
 ## - Model page: <a href='https://huggingface.co/{username}/{slugged_lora_name}'>{username}/{slugged_lora_name}</a> <small>(will be available when training finishes)</small>"""
-    else:
-        print("Error: ", outcome.stderr)
-        raise gr.Error("Something went wrong. Make sure the name of your LoRA is unique and try again")
 
 def calculate_price(iterations, with_prior_preservation):
     if(with_prior_preservation):
@@ -636,24 +693,16 @@ 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):
+def _run_captioning(*inputs):
     images = inputs[0]
     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):
         concept_caption = inputs[index + 1].strip()
         try:
-            generated_text = caption_with_mage(caption_client, image, caption_instruction)
+            generated_text = caption_with_mage(image, caption_instruction)
         except Exception as exc:
             raise gr.Error(
                 f"Mage-VL failed while captioning image {index + 1}: {exc}"
@@ -666,6 +715,19 @@ def run_captioning(*inputs):
         yield final_captions
 
 
+def captioning_duration(*inputs):
+    images = inputs[0] or []
+    return min(180, max(30, 20 + len(images) * 3))
+
+
+if is_spaces:
+    @spaces.GPU(duration=captioning_duration)
+    def run_captioning(*inputs):
+        yield from _run_captioning(*inputs)
+else:
+    run_captioning = _run_captioning
+
+
 def export_captions(images, *captions):
     if not images:
         raise gr.Error("Upload images before exporting captions.")

requirements.txt

@@ -1,10 +1,12 @@
-autotrain-advanced
-numpy==1.26.4
-huggingface-hub==0.27.0
+numpy<2
+huggingface-hub>=1.5.0,<2
 python-slugify
-gradio==4.44.1
-gradio_client==1.3.0
-fastapi==0.115.6
-starlette==0.41.3
-jinja2==3.1.4
+gradio==5.29.0
+gradio_client==1.10.0
 spaces
+transformers==5.7.0
+accelerate
+safetensors
+pillow
+torchvision
+opencv-python-headless<4.12

requirements_local.txt

@@ -1,13 +1,17 @@
 torch
 torchvision
+numpy<2
 python-slugify
-gradio>=4.37,<6
+gradio==5.29.0
+gradio_client==1.10.0
+huggingface-hub>=1.5.0,<2
+transformers==5.7.0
+accelerate
+safetensors
+pillow
+opencv-python-headless<4.12
 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
 datasets>=2.20.0
 ftfy