Upload 31 files

.dockerignore

@@ -0,0 +1,16 @@
+.cache/
+cudnn_windows/
+bitsandbytes_windows/
+bitsandbytes_windows_deprecated/
+dataset/
+__pycache__/
+venv/
+**/.hadolint.yml
+**/*.log
+**/.git
+**/.gitignore
+**/.env
+**/.github
+**/.vscode
+**/*.ps1
+sd-scripts/

.gitattributes

@@ -1,36 +1,37 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-flow.gif filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+sample.png filter=lfs diff=lfs merge=lfs -text
+seed.gif filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,53 @@
+# Base image with CUDA 12.2
+FROM nvidia/cuda:12.2.2-base-ubuntu22.04
+
+# Install pip if not already installed
+RUN apt-get update -y && apt-get install -y \
+    python3-pip \
+    python3-dev \
+    git \
+    build-essential  # Install dependencies for building extensions
+
+# Define environment variables for UID and GID and local timezone
+ENV PUID=${PUID:-1000}
+ENV PGID=${PGID:-1000}
+
+# Create a group with the specified GID
+RUN groupadd -g "${PGID}" appuser
+# Create a user with the specified UID and GID
+RUN useradd -m -s /bin/sh -u "${PUID}" -g "${PGID}" appuser
+
+WORKDIR /app
+
+# Get sd-scripts from kohya-ss and install them
+RUN git clone -b sd3 https://github.com/kohya-ss/sd-scripts && \
+    cd sd-scripts && \
+    pip install --no-cache-dir -r ./requirements.txt
+
+# Install main application dependencies
+COPY ./requirements.txt ./requirements.txt
+RUN pip install --no-cache-dir -r ./requirements.txt
+
+# Install Torch, Torchvision, and Torchaudio for CUDA 12.2
+RUN pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu122/torch_stable.html
+
+RUN chown -R appuser:appuser /app
+
+# delete redundant requirements.txt and sd-scripts directory within the container
+RUN rm -r ./sd-scripts
+RUN rm ./requirements.txt
+
+#Run application as non-root
+USER appuser
+
+# Copy fluxgym application code
+COPY . ./fluxgym
+
+EXPOSE 7860
+
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+
+WORKDIR /app/fluxgym
+
+# Run fluxgym Python application
+CMD ["python3", "./app.py"]

Dockerfile.cuda12.4

@@ -0,0 +1,53 @@
+# Base image with CUDA 12.4
+FROM nvidia/cuda:12.4.1-base-ubuntu22.04
+
+# Install pip if not already installed
+RUN apt-get update -y && apt-get install -y \
+    python3-pip \
+    python3-dev \
+    git \
+    build-essential  # Install dependencies for building extensions
+
+# Define environment variables for UID and GID and local timezone
+ENV PUID=${PUID:-1000}
+ENV PGID=${PGID:-1000}
+
+# Create a group with the specified GID
+RUN groupadd -g "${PGID}" appuser
+# Create a user with the specified UID and GID
+RUN useradd -m -s /bin/sh -u "${PUID}" -g "${PGID}" appuser
+
+WORKDIR /app
+
+# Get sd-scripts from kohya-ss and install them
+RUN git clone -b sd3 https://github.com/kohya-ss/sd-scripts && \
+    cd sd-scripts && \
+    pip install --no-cache-dir -r ./requirements.txt
+
+# Install main application dependencies
+COPY ./requirements.txt ./requirements.txt
+RUN pip install --no-cache-dir -r ./requirements.txt
+
+# Install Torch, Torchvision, and Torchaudio for CUDA 12.4
+RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
+
+RUN chown -R appuser:appuser /app
+
+# delete redundant requirements.txt and sd-scripts directory within the container
+RUN rm -r ./sd-scripts
+RUN rm ./requirements.txt
+
+#Run application as non-root
+USER appuser
+
+# Copy fluxgym application code
+COPY . ./fluxgym
+
+EXPOSE 7860
+
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+
+WORKDIR /app/fluxgym
+
+# Run fluxgym Python application
+CMD ["python3", "./app.py"]

LICENSE

@@ -0,0 +1,7 @@
+Copyright 2024 cocktailpeanut
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

README.md

@@ -1,11 +1,11 @@
----
-title: FLuxGym
-emoji: 🏋️
-colorFrom: green
-colorTo: blue
-sdk: docker
-app_port: 7860
-pinned: True
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: FLuxGym
+emoji: 🏋️
+colorFrom: green
+colorTo: blue
+sdk: docker
+app_port: 7860
+pinned: True
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app-launch.sh

@@ -0,0 +1,5 @@
+#!/usr/bin/env bash
+
+cd "`dirname "$0"`" || exit 1
+. env/bin/activate
+python app.py

app.py

@@ -0,0 +1,1119 @@
+import os
+import sys
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+os.environ['GRADIO_ANALYTICS_ENABLED'] = '0'
+sys.path.insert(0, os.getcwd())
+sys.path.append(os.path.join(os.path.dirname(__file__), 'sd-scripts'))
+import subprocess
+import gradio as gr
+from PIL import Image
+import torch
+import uuid
+import shutil
+import json
+import yaml
+from slugify import slugify
+from transformers import AutoProcessor, AutoModelForCausalLM
+from gradio_logsview import LogsView, LogsViewRunner
+from huggingface_hub import hf_hub_download, HfApi
+from library import flux_train_utils, huggingface_util
+from argparse import Namespace
+import train_network
+import toml
+import re
+MAX_IMAGES = 150
+
+with open('models.yaml', 'r') as file:
+    models = yaml.safe_load(file)
+
+def readme(base_model, lora_name, instance_prompt, sample_prompts):
+
+    # model license
+    model_config = models[base_model]
+    model_file = model_config["file"]
+    base_model_name = model_config["base"]
+    license = None
+    license_name = None
+    license_link = None
+    license_items = []
+    if "license" in model_config:
+        license = model_config["license"]
+        license_items.append(f"license: {license}")
+    if "license_name" in model_config:
+        license_name = model_config["license_name"]
+        license_items.append(f"license_name: {license_name}")
+    if "license_link" in model_config:
+        license_link = model_config["license_link"]
+        license_items.append(f"license_link: {license_link}")
+    license_str = "\n".join(license_items)
+    print(f"license_items={license_items}")
+    print(f"license_str = {license_str}")
+
+    # tags
+    tags = [ "text-to-image", "flux", "lora", "diffusers", "template:sd-lora", "fluxgym" ]
+
+    # widgets
+    widgets = []
+    sample_image_paths = []
+    output_name = slugify(lora_name)
+    samples_dir = resolve_path_without_quotes(f"outputs/{output_name}/sample")
+    try:
+        for filename in os.listdir(samples_dir):
+            # Filename Schema: [name]_[steps]_[index]_[timestamp].png
+            match = re.search(r"_(\d+)_(\d+)_(\d+)\.png$", filename)
+            if match:
+                steps, index, timestamp = int(match.group(1)), int(match.group(2)), int(match.group(3))
+                sample_image_paths.append((steps, index, f"sample/{filename}"))
+
+        # Sort by numeric index
+        sample_image_paths.sort(key=lambda x: x[0], reverse=True)
+
+        final_sample_image_paths = sample_image_paths[:len(sample_prompts)]
+        final_sample_image_paths.sort(key=lambda x: x[1])
+        for i, prompt in enumerate(sample_prompts):
+            _, _, image_path = final_sample_image_paths[i]
+            widgets.append(
+                {
+                    "text": prompt,
+                    "output": {
+                        "url": image_path
+                    },
+                }
+            )
+    except:
+        print(f"no samples")
+    dtype = "torch.bfloat16"
+    # Construct the README content
+    readme_content = f"""---
+tags:
+{yaml.dump(tags, indent=4).strip()}
+{"widget:" if os.path.isdir(samples_dir) else ""}
+{yaml.dump(widgets, indent=4).strip() if widgets else ""}
+base_model: {base_model_name}
+{"instance_prompt: " + instance_prompt if instance_prompt else ""}
+{license_str}
+---
+
+# {lora_name}
+
+A Flux LoRA trained on a local computer with [Fluxgym](https://github.com/cocktailpeanut/fluxgym)
+
+<Gallery />
+
+## Trigger words
+
+{"You should use `" + instance_prompt + "` to trigger the image generation." if instance_prompt else "No trigger words defined."}
+
+## Download model and use it with ComfyUI, AUTOMATIC1111, SD.Next, Invoke AI, Forge, etc.
+
+Weights for this model are available in Safetensors format.
+
+"""
+    return readme_content
+
+def account_hf():
+    try:
+        with open("HF_TOKEN", "r") as file:
+            token = file.read()
+            api = HfApi(token=token)
+            try:
+                account = api.whoami()
+                return { "token": token, "account": account['name'] }
+            except:
+                return None
+    except:
+        return None
+
+"""
+hf_logout.click(fn=logout_hf, outputs=[hf_token, hf_login, hf_logout, repo_owner])
+"""
+def logout_hf():
+    os.remove("HF_TOKEN")
+    global current_account
+    current_account = account_hf()
+    print(f"current_account={current_account}")
+    return gr.update(value=""), gr.update(visible=True), gr.update(visible=False), gr.update(value="", visible=False)
+
+
+"""
+hf_login.click(fn=login_hf, inputs=[hf_token], outputs=[hf_token, hf_login, hf_logout, repo_owner])
+"""
+def login_hf(hf_token):
+    api = HfApi(token=hf_token)
+    try:
+        account = api.whoami()
+        if account != None:
+            if "name" in account:
+                with open("HF_TOKEN", "w") as file:
+                    file.write(hf_token)
+                global current_account
+                current_account = account_hf()
+                return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(value=current_account["account"], visible=True)
+        return gr.update(), gr.update(), gr.update(), gr.update()
+    except:
+        print(f"incorrect hf_token")
+        return gr.update(), gr.update(), gr.update(), gr.update()
+
+def upload_hf(base_model, lora_rows, repo_owner, repo_name, repo_visibility, hf_token):
+    src = lora_rows
+    repo_id = f"{repo_owner}/{repo_name}"
+    gr.Info(f"Uploading to Huggingface. Please Stand by...", duration=None)
+    args = Namespace(
+        huggingface_repo_id=repo_id,
+        huggingface_repo_type="model",
+        huggingface_repo_visibility=repo_visibility,
+        huggingface_path_in_repo="",
+        huggingface_token=hf_token,
+        async_upload=False
+    )
+    print(f"upload_hf args={args}")
+    huggingface_util.upload(args=args, src=src)
+    gr.Info(f"[Upload Complete] https://huggingface.co/{repo_id}", duration=None)
+
+def load_captioning(uploaded_files, concept_sentence):
+    uploaded_images = [file for file in uploaded_files if not file.endswith('.txt')]
+    txt_files = [file for file in uploaded_files if file.endswith('.txt')]
+    txt_files_dict = {os.path.splitext(os.path.basename(txt_file))[0]: txt_file for txt_file in txt_files}
+    updates = []
+    if len(uploaded_images) <= 1:
+        raise gr.Error(
+            "Please upload at least 2 images to train your model (the ideal number with default settings is between 4-30)"
+        )
+    elif len(uploaded_images) > MAX_IMAGES:
+        raise gr.Error(f"For now, only {MAX_IMAGES} or less images are allowed for training")
+    # Update for the captioning_area
+    # for _ in range(3):
+    updates.append(gr.update(visible=True))
+    # Update visibility and image for each captioning row and image
+    for i in range(1, MAX_IMAGES + 1):
+        # Determine if the current row and image should be visible
+        visible = i <= len(uploaded_images)
+
+        # Update visibility of the captioning row
+        updates.append(gr.update(visible=visible))
+
+        # Update for image component - display image if available, otherwise hide
+        image_value = uploaded_images[i - 1] if visible else None
+        updates.append(gr.update(value=image_value, visible=visible))
+
+        corresponding_caption = False
+        if(image_value):
+            base_name = os.path.splitext(os.path.basename(image_value))[0]
+            if base_name in txt_files_dict:
+                with open(txt_files_dict[base_name], 'r') as file:
+                    corresponding_caption = file.read()
+
+        # Update value of captioning area
+        text_value = corresponding_caption if visible and corresponding_caption else concept_sentence if visible and concept_sentence else None
+        updates.append(gr.update(value=text_value, visible=visible))
+
+    # Update for the sample caption area
+    updates.append(gr.update(visible=True))
+    updates.append(gr.update(visible=True))
+
+    return updates
+
+def hide_captioning():
+    return gr.update(visible=False), gr.update(visible=False)
+
+def resize_image(image_path, output_path, size):
+    with Image.open(image_path) as img:
+        width, height = img.size
+        if width < height:
+            new_width = size
+            new_height = int((size/width) * height)
+        else:
+            new_height = size
+            new_width = int((size/height) * width)
+        print(f"resize {image_path} : {new_width}x{new_height}")
+        img_resized = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
+        img_resized.save(output_path)
+
+def create_dataset(destination_folder, size, *inputs):
+    print("Creating dataset")
+    images = inputs[0]
+    if not os.path.exists(destination_folder):
+        os.makedirs(destination_folder)
+
+    for index, image in enumerate(images):
+        # copy the images to the datasets folder
+        new_image_path = shutil.copy(image, destination_folder)
+
+        # if it's a caption text file skip the next bit
+        ext = os.path.splitext(new_image_path)[-1].lower()
+        if ext == '.txt':
+            continue
+
+        # resize the images
+        resize_image(new_image_path, new_image_path, size)
+
+        # copy the captions
+
+        original_caption = inputs[index + 1]
+
+        image_file_name = os.path.basename(new_image_path)
+        caption_file_name = os.path.splitext(image_file_name)[0] + ".txt"
+        caption_path = resolve_path_without_quotes(os.path.join(destination_folder, caption_file_name))
+        print(f"image_path={new_image_path}, caption_path = {caption_path}, original_caption={original_caption}")
+        # if caption_path exists, do not write
+        if os.path.exists(caption_path):
+            print(f"{caption_path} already exists. use the existing .txt file")
+        else:
+            print(f"{caption_path} create a .txt caption file")
+            with open(caption_path, 'w') as file:
+                file.write(original_caption)
+
+    print(f"destination_folder {destination_folder}")
+    return destination_folder
+
+
+def run_captioning(images, concept_sentence, *captions):
+    print(f"run_captioning")
+    print(f"concept sentence {concept_sentence}")
+    print(f"captions {captions}")
+    #Load internally to not consume resources for training
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"device={device}")
+    torch_dtype = torch.float16
+    model = AutoModelForCausalLM.from_pretrained(
+        "multimodalart/Florence-2-large-no-flash-attn", torch_dtype=torch_dtype, trust_remote_code=True
+    ).to(device)
+    processor = AutoProcessor.from_pretrained("multimodalart/Florence-2-large-no-flash-attn", trust_remote_code=True)
+
+    captions = list(captions)
+    for i, image_path in enumerate(images):
+        print(captions[i])
+        if isinstance(image_path, str):  # If image is a file path
+            image = Image.open(image_path).convert("RGB")
+
+        prompt = "<DETAILED_CAPTION>"
+        inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
+        print(f"inputs {inputs}")
+
+        generated_ids = model.generate(
+            input_ids=inputs["input_ids"], pixel_values=inputs["pixel_values"], max_new_tokens=1024, num_beams=3
+        )
+        print(f"generated_ids {generated_ids}")
+
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        print(f"generated_text: {generated_text}")
+        parsed_answer = processor.post_process_generation(
+            generated_text, task=prompt, image_size=(image.width, image.height)
+        )
+        print(f"parsed_answer = {parsed_answer}")
+        caption_text = parsed_answer["<DETAILED_CAPTION>"].replace("The image shows ", "")
+        print(f"caption_text = {caption_text}, concept_sentence={concept_sentence}")
+        if concept_sentence:
+            caption_text = f"{concept_sentence} {caption_text}"
+        captions[i] = caption_text
+
+        yield captions
+    model.to("cpu")
+    del model
+    del processor
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+def recursive_update(d, u):
+    for k, v in u.items():
+        if isinstance(v, dict) and v:
+            d[k] = recursive_update(d.get(k, {}), v)
+        else:
+            d[k] = v
+    return d
+
+def download(base_model):
+    model = models[base_model]
+    model_file = model["file"]
+    repo = model["repo"]
+
+    # download unet
+    if base_model == "flux-dev" or base_model == "flux-schnell":
+        unet_folder = "models/unet"
+    else:
+        unet_folder = f"models/unet/{repo}"
+    unet_path = os.path.join(unet_folder, model_file)
+    if not os.path.exists(unet_path):
+        os.makedirs(unet_folder, exist_ok=True)
+        gr.Info(f"Downloading base model: {base_model}. Please wait. (You can check the terminal for the download progress)", duration=None)
+        print(f"download {base_model}")
+        hf_hub_download(repo_id=repo, local_dir=unet_folder, filename=model_file)
+
+    # download vae
+    vae_folder = "models/vae"
+    vae_path = os.path.join(vae_folder, "ae.sft")
+    if not os.path.exists(vae_path):
+        os.makedirs(vae_folder, exist_ok=True)
+        gr.Info(f"Downloading vae")
+        print(f"downloading ae.sft...")
+        hf_hub_download(repo_id="cocktailpeanut/xulf-dev", local_dir=vae_folder, filename="ae.sft")
+
+    # download clip
+    clip_folder = "models/clip"
+    clip_l_path = os.path.join(clip_folder, "clip_l.safetensors")
+    if not os.path.exists(clip_l_path):
+        os.makedirs(clip_folder, exist_ok=True)
+        gr.Info(f"Downloading clip...")
+        print(f"download clip_l.safetensors")
+        hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", local_dir=clip_folder, filename="clip_l.safetensors")
+
+    # download t5xxl
+    t5xxl_path = os.path.join(clip_folder, "t5xxl_fp16.safetensors")
+    if not os.path.exists(t5xxl_path):
+        print(f"download t5xxl_fp16.safetensors")
+        gr.Info(f"Downloading t5xxl...")
+        hf_hub_download(repo_id="comfyanonymous/flux_text_encoders", local_dir=clip_folder, filename="t5xxl_fp16.safetensors")
+
+
+def resolve_path(p):
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    norm_path = os.path.normpath(os.path.join(current_dir, p))
+    return f"\"{norm_path}\""
+def resolve_path_without_quotes(p):
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    norm_path = os.path.normpath(os.path.join(current_dir, p))
+    return norm_path
+
+def gen_sh(
+    base_model,
+    output_name,
+    resolution,
+    seed,
+    workers,
+    learning_rate,
+    network_dim,
+    max_train_epochs,
+    save_every_n_epochs,
+    timestep_sampling,
+    guidance_scale,
+    vram,
+    sample_prompts,
+    sample_every_n_steps,
+    *advanced_components
+):
+
+    print(f"gen_sh: network_dim:{network_dim}, max_train_epochs={max_train_epochs}, save_every_n_epochs={save_every_n_epochs}, timestep_sampling={timestep_sampling}, guidance_scale={guidance_scale}, vram={vram}, sample_prompts={sample_prompts}, sample_every_n_steps={sample_every_n_steps}")
+
+    output_dir = resolve_path(f"outputs/{output_name}")
+    sample_prompts_path = resolve_path(f"outputs/{output_name}/sample_prompts.txt")
+
+    line_break = "\\"
+    file_type = "sh"
+    if sys.platform == "win32":
+        line_break = "^"
+        file_type = "bat"
+
+    ############# Sample args ########################
+    sample = ""
+    if len(sample_prompts) > 0 and sample_every_n_steps > 0:
+        sample = f"""--sample_prompts={sample_prompts_path} --sample_every_n_steps="{sample_every_n_steps}" {line_break}"""
+
+
+    ############# Optimizer args ########################
+#    if vram == "8G":
+#        optimizer = f"""--optimizer_type adafactor {line_break}
+#    --optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" {line_break}
+#        --split_mode {line_break}
+#        --network_args "train_blocks=single" {line_break}
+#        --lr_scheduler constant_with_warmup {line_break}
+#        --max_grad_norm 0.0 {line_break}"""
+    if vram == "16G":
+        # 16G VRAM
+        optimizer = f"""--optimizer_type adafactor {line_break}
+  --optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" {line_break}
+  --lr_scheduler constant_with_warmup {line_break}
+  --max_grad_norm 0.0 {line_break}"""
+    elif vram == "12G":
+      # 12G VRAM
+        optimizer = f"""--optimizer_type adafactor {line_break}
+  --optimizer_args "relative_step=False" "scale_parameter=False" "warmup_init=False" {line_break}
+  --split_mode {line_break}
+  --network_args "train_blocks=single" {line_break}
+  --lr_scheduler constant_with_warmup {line_break}
+  --max_grad_norm 0.0 {line_break}"""
+    else:
+        # 20G+ VRAM
+        optimizer = f"--optimizer_type adamw8bit {line_break}"
+
+
+    #######################################################
+    model_config = models[base_model]
+    model_file = model_config["file"]
+    repo = model_config["repo"]
+    if base_model == "flux-dev" or base_model == "flux-schnell":
+        model_folder = "models/unet"
+    else:
+        model_folder = f"models/unet/{repo}"
+    model_path = os.path.join(model_folder, model_file)
+    pretrained_model_path = resolve_path(model_path)
+
+    clip_path = resolve_path("models/clip/clip_l.safetensors")
+    t5_path = resolve_path("models/clip/t5xxl_fp16.safetensors")
+    ae_path = resolve_path("models/vae/ae.sft")
+    sh = f"""accelerate launch {line_break}
+  --mixed_precision bf16 {line_break}
+  --num_cpu_threads_per_process 1 {line_break}
+  sd-scripts/flux_train_network.py {line_break}
+  --pretrained_model_name_or_path {pretrained_model_path} {line_break}
+  --clip_l {clip_path} {line_break}
+  --t5xxl {t5_path} {line_break}
+  --ae {ae_path} {line_break}
+  --cache_latents_to_disk {line_break}
+  --save_model_as safetensors {line_break}
+  --sdpa --persistent_data_loader_workers {line_break}
+  --max_data_loader_n_workers {workers} {line_break}
+  --seed {seed} {line_break}
+  --gradient_checkpointing {line_break}
+  --mixed_precision bf16 {line_break}
+  --save_precision bf16 {line_break}
+  --network_module networks.lora_flux {line_break}
+  --network_dim {network_dim} {line_break}
+  {optimizer}{sample}
+  --learning_rate {learning_rate} {line_break}
+  --cache_text_encoder_outputs {line_break}
+  --cache_text_encoder_outputs_to_disk {line_break}
+  --fp8_base {line_break}
+  --highvram {line_break}
+  --max_train_epochs {max_train_epochs} {line_break}
+  --save_every_n_epochs {save_every_n_epochs} {line_break}
+  --dataset_config {resolve_path(f"outputs/{output_name}/dataset.toml")} {line_break}
+  --output_dir {output_dir} {line_break}
+  --output_name {output_name} {line_break}
+  --timestep_sampling {timestep_sampling} {line_break}
+  --discrete_flow_shift 3.1582 {line_break}
+  --model_prediction_type raw {line_break}
+  --guidance_scale {guidance_scale} {line_break}
+  --loss_type l2 {line_break}"""
+   
+
+
+    ############# Advanced args ########################
+    global advanced_component_ids
+    global original_advanced_component_values
+   
+    # check dirty
+    print(f"original_advanced_component_values = {original_advanced_component_values}")
+    advanced_flags = []
+    for i, current_value in enumerate(advanced_components):
+#        print(f"compare {advanced_component_ids[i]}: old={original_advanced_component_values[i]}, new={current_value}")
+        if original_advanced_component_values[i] != current_value:
+            # dirty
+            if current_value == True:
+                # Boolean
+                advanced_flags.append(advanced_component_ids[i])
+            else:
+                # string
+                advanced_flags.append(f"{advanced_component_ids[i]} {current_value}")
+
+    if len(advanced_flags) > 0:
+        advanced_flags_str = f" {line_break}\n  ".join(advanced_flags)
+        sh = sh + "\n  " + advanced_flags_str
+
+    return sh
+
+def gen_toml(
+  dataset_folder,
+  resolution,
+  class_tokens,
+  num_repeats
+):
+    toml = f"""[general]
+shuffle_caption = false
+caption_extension = '.txt'
+keep_tokens = 1
+
+[[datasets]]
+resolution = {resolution}
+batch_size = 1
+keep_tokens = 1
+
+  [[datasets.subsets]]
+  image_dir = '{resolve_path_without_quotes(dataset_folder)}'
+  class_tokens = '{class_tokens}'
+  num_repeats = {num_repeats}"""
+    return toml
+
+def update_total_steps(max_train_epochs, num_repeats, images):
+    try:
+        num_images = len(images)
+        total_steps = max_train_epochs * num_images * num_repeats
+        print(f"max_train_epochs={max_train_epochs} num_images={num_images}, num_repeats={num_repeats}, total_steps={total_steps}")
+        return gr.update(value = total_steps)
+    except:
+        print("")
+
+def set_repo(lora_rows):
+    selected_name = os.path.basename(lora_rows)
+    return gr.update(value=selected_name)
+
+def get_loras():
+    try:
+        outputs_path = resolve_path_without_quotes(f"outputs")
+        files = os.listdir(outputs_path)
+        folders = [os.path.join(outputs_path, item) for item in files if os.path.isdir(os.path.join(outputs_path, item)) and item != "sample"]
+        folders.sort(key=lambda file: os.path.getctime(file), reverse=True)
+        return folders
+    except Exception as e:
+        return []
+
+def get_samples(lora_name):
+    output_name = slugify(lora_name)
+    try:
+        samples_path = resolve_path_without_quotes(f"outputs/{output_name}/sample")
+        files = [os.path.join(samples_path, file) for file in os.listdir(samples_path)]
+        files.sort(key=lambda file: os.path.getctime(file), reverse=True)
+        return files
+    except:
+        return []
+
+def start_training(
+    base_model,
+    lora_name,
+    train_script,
+    train_config,
+    sample_prompts,
+):
+    # write custom script and toml
+    if not os.path.exists("models"):
+        os.makedirs("models", exist_ok=True)
+    if not os.path.exists("outputs"):
+        os.makedirs("outputs", exist_ok=True)
+    output_name = slugify(lora_name)
+    output_dir = resolve_path_without_quotes(f"outputs/{output_name}")
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir, exist_ok=True)
+
+    download(base_model)
+
+    file_type = "sh"
+    if sys.platform == "win32":
+        file_type = "bat"
+
+    sh_filename = f"train.{file_type}"
+    sh_filepath = resolve_path_without_quotes(f"outputs/{output_name}/{sh_filename}")
+    with open(sh_filepath, 'w', encoding="utf-8") as file:
+        file.write(train_script)
+    gr.Info(f"Generated train script at {sh_filename}")
+
+
+    dataset_path = resolve_path_without_quotes(f"outputs/{output_name}/dataset.toml")
+    with open(dataset_path, 'w', encoding="utf-8") as file:
+        file.write(train_config)
+    gr.Info(f"Generated dataset.toml")
+
+    sample_prompts_path = resolve_path_without_quotes(f"outputs/{output_name}/sample_prompts.txt")
+    with open(sample_prompts_path, 'w', encoding='utf-8') as file:
+        file.write(sample_prompts)
+    gr.Info(f"Generated sample_prompts.txt")
+
+    # Train
+    if sys.platform == "win32":
+        command = sh_filepath
+    else:
+        command = f"bash \"{sh_filepath}\""
+
+    # Use Popen to run the command and capture output in real-time
+    env = os.environ.copy()
+    env['PYTHONIOENCODING'] = 'utf-8'
+    env['LOG_LEVEL'] = 'DEBUG'
+    runner = LogsViewRunner()
+    cwd = os.path.dirname(os.path.abspath(__file__))
+    gr.Info(f"Started training")
+    yield from runner.run_command([command], cwd=cwd)
+    yield runner.log(f"Runner: {runner}")
+
+    # Generate Readme
+    config = toml.loads(train_config)
+    concept_sentence = config['datasets'][0]['subsets'][0]['class_tokens']
+    print(f"concept_sentence={concept_sentence}")
+    print(f"lora_name {lora_name}, concept_sentence={concept_sentence}, output_name={output_name}")
+    sample_prompts_path = resolve_path_without_quotes(f"outputs/{output_name}/sample_prompts.txt")
+    with open(sample_prompts_path, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+    sample_prompts = [line.strip() for line in lines if len(line.strip()) > 0 and line[0] != "#"]
+    md = readme(base_model, lora_name, concept_sentence, sample_prompts)
+    readme_path = resolve_path_without_quotes(f"outputs/{output_name}/README.md")
+    with open(readme_path, "w", encoding="utf-8") as f:
+        f.write(md)
+
+    gr.Info(f"Training Complete. Check the outputs folder for the LoRA files.", duration=None)
+
+
+def update(
+    base_model,
+    lora_name,
+    resolution,
+    seed,
+    workers,
+    class_tokens,
+    learning_rate,
+    network_dim,
+    max_train_epochs,
+    save_every_n_epochs,
+    timestep_sampling,
+    guidance_scale,
+    vram,
+    num_repeats,
+    sample_prompts,
+    sample_every_n_steps,
+    *advanced_components,
+):
+    output_name = slugify(lora_name)
+    dataset_folder = str(f"datasets/{output_name}")
+    sh = gen_sh(
+        base_model,
+        output_name,
+        resolution,
+        seed,
+        workers,
+        learning_rate,
+        network_dim,
+        max_train_epochs,
+        save_every_n_epochs,
+        timestep_sampling,
+        guidance_scale,
+        vram,
+        sample_prompts,
+        sample_every_n_steps,
+        *advanced_components,
+    )
+    toml = gen_toml(
+        dataset_folder,
+        resolution,
+        class_tokens,
+        num_repeats
+    )
+    return gr.update(value=sh), gr.update(value=toml), dataset_folder
+
+"""
+demo.load(fn=loaded, js=js, outputs=[hf_token, hf_login, hf_logout, hf_account])
+"""
+def loaded():
+    global current_account
+    current_account = account_hf()
+    print(f"current_account={current_account}")
+    if current_account != None:
+        return gr.update(value=current_account["token"]), gr.update(visible=False), gr.update(visible=True), gr.update(value=current_account["account"], visible=True)
+    else:
+        return gr.update(value=""), gr.update(visible=True), gr.update(visible=False), gr.update(value="", visible=False)
+
+def update_sample(concept_sentence):
+    return gr.update(value=concept_sentence)
+
+def refresh_publish_tab():
+    loras = get_loras()
+    return gr.Dropdown(label="Trained LoRAs", choices=loras)
+
+def init_advanced():
+    # if basic_args
+    basic_args = {
+        'pretrained_model_name_or_path',
+        'clip_l',
+        't5xxl',
+        'ae',
+        'cache_latents_to_disk',
+        'save_model_as',
+        'sdpa',
+        'persistent_data_loader_workers',
+        'max_data_loader_n_workers',
+        'seed',
+        'gradient_checkpointing',
+        'mixed_precision',
+        'save_precision',
+        'network_module',
+        'network_dim',
+        'learning_rate',
+        'cache_text_encoder_outputs',
+        'cache_text_encoder_outputs_to_disk',
+        'fp8_base',
+        'highvram',
+        'max_train_epochs',
+        'save_every_n_epochs',
+        'dataset_config',
+        'output_dir',
+        'output_name',
+        'timestep_sampling',
+        'discrete_flow_shift',
+        'model_prediction_type',
+        'guidance_scale',
+        'loss_type',
+        'optimizer_type',
+        'optimizer_args',
+        'lr_scheduler',
+        'sample_prompts',
+        'sample_every_n_steps',
+        'max_grad_norm',
+        'split_mode',
+        'network_args'
+    }
+
+    # generate a UI config
+    # if not in basic_args, create a simple form
+    parser = train_network.setup_parser()
+    flux_train_utils.add_flux_train_arguments(parser)
+    args_info = {}
+    for action in parser._actions:
+        if action.dest != 'help':  # Skip the default help argument
+            # if the dest is included in basic_args
+            args_info[action.dest] = {
+                "action": action.option_strings,  # Option strings like '--use_8bit_adam'
+                "type": action.type,              # Type of the argument
+                "help": action.help,              # Help message
+                "default": action.default,        # Default value, if any
+                "required": action.required       # Whether the argument is required
+            }
+    temp = []
+    for key in args_info:
+        temp.append({ 'key': key, 'action': args_info[key] })
+    temp.sort(key=lambda x: x['key'])
+    advanced_component_ids = []
+    advanced_components = []
+    for item in temp:
+        key = item['key']
+        action = item['action']
+        if key in basic_args:
+            print("")
+        else:
+            action_type = str(action['type'])
+            component = None
+            with gr.Column(min_width=300):
+                if action_type == "None":
+                    # radio
+                    component = gr.Checkbox()
+    #            elif action_type == "<class 'str'>":
+    #                component = gr.Textbox()
+    #            elif action_type == "<class 'int'>":
+    #                component = gr.Number(precision=0)
+    #            elif action_type == "<class 'float'>":
+    #                component = gr.Number()
+    #            elif "int_or_float" in action_type:
+    #                component = gr.Number()
+                else:
+                    component = gr.Textbox(value="")
+                if component != None:
+                    component.interactive = True
+                    component.elem_id = action['action'][0]
+                    component.label = component.elem_id
+                    component.elem_classes = ["advanced"]
+                if action['help'] != None:
+                    component.info = action['help']
+            advanced_components.append(component)
+            advanced_component_ids.append(component.elem_id)
+    return advanced_components, advanced_component_ids
+
+
+theme = gr.themes.Monochrome(
+    text_size=gr.themes.Size(lg="18px", md="15px", sm="13px", xl="22px", xs="12px", xxl="24px", xxs="9px"),
+    font=[gr.themes.GoogleFont("Source Sans Pro"), "ui-sans-serif", "system-ui", "sans-serif"],
+)
+css = """
+@keyframes rotate {
+    0% {
+        transform: rotate(0deg);
+    }
+    100% {
+        transform: rotate(360deg);
+    }
+}
+#advanced_options .advanced:nth-child(even) { background: rgba(0,0,100,0.04) !important; }
+h1{font-family: georgia; font-style: italic; font-weight: bold; font-size: 30px; letter-spacing: -1px;}
+h3{margin-top: 0}
+.tabitem{border: 0px}
+.group_padding{}
+nav{position: fixed; top: 0; left: 0; right: 0; z-index: 1000; text-align: center; padding: 10px; box-sizing: border-box; display: flex; align-items: center; backdrop-filter: blur(10px); }
+nav button { background: none; color: firebrick; font-weight: bold; border: 2px solid firebrick; padding: 5px 10px; border-radius: 5px; font-size: 14px; }
+nav img { height: 40px; width: 40px; border-radius: 40px; }
+nav img.rotate { animation: rotate 2s linear infinite; }
+.flexible { flex-grow: 1; }
+.tast-details { margin: 10px 0 !important; }
+.toast-wrap { bottom: var(--size-4) !important; top: auto !important; border: none !important; backdrop-filter: blur(10px); }
+.toast-title, .toast-text, .toast-icon, .toast-close { color: black !important; font-size: 14px; }
+.toast-body { border: none !important; }
+#terminal { box-shadow: none !important; margin-bottom: 25px; background: rgba(0,0,0,0.03); }
+#terminal .generating { border: none !important; }
+#terminal label { position: absolute !important; }
+.tabs { margin-top: 50px; }
+.hidden { display: none !important; }
+.codemirror-wrapper .cm-line { font-size: 12px !important; }
+label { font-weight: bold !important; }
+#start_training.clicked { background: silver; color: black; }
+"""
+
+js = """
+function() {
+    let autoscroll = document.querySelector("#autoscroll")
+    if (window.iidxx) {
+        window.clearInterval(window.iidxx);
+    }
+    window.iidxx = window.setInterval(function() {
+        let text=document.querySelector(".codemirror-wrapper .cm-line").innerText.trim()
+        let img = document.querySelector("#logo")
+        if (text.length > 0) {
+            autoscroll.classList.remove("hidden")
+            if (autoscroll.classList.contains("on")) {
+                autoscroll.textContent = "Autoscroll ON"
+                window.scrollTo(0, document.body.scrollHeight, { behavior: "smooth" });
+                img.classList.add("rotate")
+            } else {
+                autoscroll.textContent = "Autoscroll OFF"
+                img.classList.remove("rotate")
+            }
+        }
+    }, 500);
+    console.log("autoscroll", autoscroll)
+    autoscroll.addEventListener("click", (e) => {
+        autoscroll.classList.toggle("on")
+    })
+    function debounce(fn, delay) {
+        let timeoutId;
+        return function(...args) {
+            clearTimeout(timeoutId);
+            timeoutId = setTimeout(() => fn(...args), delay);
+        };
+    }
+
+    function handleClick() {
+        console.log("refresh")
+        document.querySelector("#refresh").click();
+    }
+    const debouncedClick = debounce(handleClick, 1000);
+    document.addEventListener("input", debouncedClick);
+
+    document.querySelector("#start_training").addEventListener("click", (e) => {
+      e.target.classList.add("clicked")
+      e.target.innerHTML = "Training..."
+    })
+
+}
+"""
+
+current_account = account_hf()
+print(f"current_account={current_account}")
+
+with gr.Blocks(elem_id="app", theme=theme, css=css, fill_width=True) as demo:
+    with gr.Tabs() as tabs:
+        with gr.TabItem("Gym"):
+            output_components = []
+            with gr.Row():
+                gr.HTML("""<nav>
+            <img id='logo' src='/file=icon.png' width='80' height='80'>
+            <div class='flexible'></div>
+            <button id='autoscroll' class='on hidden'></button>
+        </nav>
+        """)
+            with gr.Row(elem_id='container'):
+                with gr.Column():
+                    gr.Markdown(
+                        """# Step 1. LoRA Info
+        <p style="margin-top:0">Configure your LoRA train settings.</p>
+        """, elem_classes="group_padding")
+                    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",
+                    )
+                    concept_sentence = gr.Textbox(
+                        elem_id="--concept_sentence",
+                        label="Trigger word/sentence",
+                        info="Trigger word or sentence to be used",
+                        placeholder="uncommon word like p3rs0n or trtcrd, or sentence like 'in the style of CNSTLL'",
+                        interactive=True,
+                    )
+                    model_names = list(models.keys())
+                    print(f"model_names={model_names}")
+                    base_model = gr.Dropdown(label="Base model (edit the models.yaml file to add more to this list)", choices=model_names, value=model_names[0])
+                    vram = gr.Radio(["20G", "16G", "12G" ], value="20G", label="VRAM", interactive=True)
+                    num_repeats = gr.Number(value=10, precision=0, label="Repeat trains per image", interactive=True)
+                    max_train_epochs = gr.Number(label="Max Train Epochs", value=16, interactive=True)
+                    total_steps = gr.Number(0, interactive=False, label="Expected training steps")
+                    sample_prompts = gr.Textbox("", lines=5, label="Sample Image Prompts (Separate with new lines)", interactive=True)
+                    sample_every_n_steps = gr.Number(0, precision=0, label="Sample Image Every N Steps", interactive=True)
+                    resolution = gr.Number(value=512, precision=0, label="Resize dataset images")
+                with gr.Column():
+                    gr.Markdown(
+                        """# Step 2. Dataset
+        <p style="margin-top:0">Make sure the captions include the trigger word.</p>
+        """, elem_classes="group_padding")
+                    with gr.Group():
+                        images = gr.File(
+                            file_types=["image", ".txt"],
+                            label="Upload your images",
+                            #info="If you want, you can also manually upload caption files that match the image names (example: img0.png => img0.txt)",
+                            file_count="multiple",
+                            interactive=True,
+                            visible=True,
+                            scale=1,
+                        )
+                    with gr.Group(visible=False) as captioning_area:
+                        do_captioning = gr.Button("Add AI captions with Florence-2")
+                        output_components.append(captioning_area)
+                        #output_components = [captioning_area]
+                        caption_list = []
+                        for i in range(1, MAX_IMAGES + 1):
+                            locals()[f"captioning_row_{i}"] = gr.Row(visible=False)
+                            with locals()[f"captioning_row_{i}"]:
+                                locals()[f"image_{i}"] = gr.Image(
+                                    type="filepath",
+                                    width=111,
+                                    height=111,
+                                    min_width=111,
+                                    interactive=False,
+                                    scale=2,
+                                    show_label=False,
+                                    show_share_button=False,
+                                    show_download_button=False,
+                                )
+                                locals()[f"caption_{i}"] = gr.Textbox(
+                                    label=f"Caption {i}", scale=15, interactive=True
+                                )
+
+                            output_components.append(locals()[f"captioning_row_{i}"])
+                            output_components.append(locals()[f"image_{i}"])
+                            output_components.append(locals()[f"caption_{i}"])
+                            caption_list.append(locals()[f"caption_{i}"])
+                with gr.Column():
+                    gr.Markdown(
+                        """# Step 3. Train
+        <p style="margin-top:0">Press start to start training.</p>
+        """, elem_classes="group_padding")
+                    refresh = gr.Button("Refresh", elem_id="refresh", visible=False)
+                    start = gr.Button("Start training", visible=False, elem_id="start_training")
+                    output_components.append(start)
+                    train_script = gr.Textbox(label="Train script", max_lines=100, interactive=True)
+                    train_config = gr.Textbox(label="Train config", max_lines=100, interactive=True)
+            with gr.Accordion("Advanced options", elem_id='advanced_options', open=False):
+                with gr.Row():
+                    with gr.Column(min_width=300):
+                        seed = gr.Number(label="--seed", info="Seed", value=42, interactive=True)
+                    with gr.Column(min_width=300):
+                        workers = gr.Number(label="--max_data_loader_n_workers", info="Number of Workers", value=2, interactive=True)
+                    with gr.Column(min_width=300):
+                        learning_rate = gr.Textbox(label="--learning_rate", info="Learning Rate", value="8e-4", interactive=True)
+                    with gr.Column(min_width=300):
+                        save_every_n_epochs = gr.Number(label="--save_every_n_epochs", info="Save every N epochs", value=4, interactive=True)
+                    with gr.Column(min_width=300):
+                        guidance_scale = gr.Number(label="--guidance_scale", info="Guidance Scale", value=1.0, interactive=True)
+                    with gr.Column(min_width=300):
+                        timestep_sampling = gr.Textbox(label="--timestep_sampling", info="Timestep Sampling", value="shift", interactive=True)
+                    with gr.Column(min_width=300):
+                        network_dim = gr.Number(label="--network_dim", info="LoRA Rank", value=4, minimum=4, maximum=128, step=4, interactive=True)
+                    advanced_components, advanced_component_ids = init_advanced()
+            with gr.Row():
+                terminal = LogsView(label="Train log", elem_id="terminal")
+            with gr.Row():
+                gallery = gr.Gallery(get_samples, inputs=[lora_name], label="Samples", every=10, columns=6)
+
+        with gr.TabItem("Publish") as publish_tab:
+            hf_token = gr.Textbox(label="Huggingface Token")
+            hf_login = gr.Button("Login")
+            hf_logout = gr.Button("Logout")
+            with gr.Row() as row:
+                gr.Markdown("**LoRA**")
+                gr.Markdown("**Upload**")
+            loras = get_loras()
+            with gr.Row():
+                lora_rows = refresh_publish_tab()
+                with gr.Column():
+                    with gr.Row():
+                        repo_owner = gr.Textbox(label="Account", interactive=False)
+                        repo_name = gr.Textbox(label="Repository Name")
+                    repo_visibility = gr.Textbox(label="Repository Visibility ('public' or 'private')", value="public")
+                    upload_button = gr.Button("Upload to HuggingFace")
+                    upload_button.click(
+                        fn=upload_hf,
+                        inputs=[
+                            base_model,
+                            lora_rows,
+                            repo_owner,
+                            repo_name,
+                            repo_visibility,
+                            hf_token,
+                        ]
+                    )
+            hf_login.click(fn=login_hf, inputs=[hf_token], outputs=[hf_token, hf_login, hf_logout, repo_owner])
+            hf_logout.click(fn=logout_hf, outputs=[hf_token, hf_login, hf_logout, repo_owner])
+
+
+    publish_tab.select(refresh_publish_tab, outputs=lora_rows)
+    lora_rows.select(fn=set_repo, inputs=[lora_rows], outputs=[repo_name])
+
+    dataset_folder = gr.State()
+
+    listeners = [
+        base_model,
+        lora_name,
+        resolution,
+        seed,
+        workers,
+        concept_sentence,
+        learning_rate,
+        network_dim,
+        max_train_epochs,
+        save_every_n_epochs,
+        timestep_sampling,
+        guidance_scale,
+        vram,
+        num_repeats,
+        sample_prompts,
+        sample_every_n_steps,
+        *advanced_components
+    ]
+    advanced_component_ids = [x.elem_id for x in advanced_components]
+    original_advanced_component_values = [comp.value for comp in advanced_components]
+    images.upload(
+        load_captioning,
+        inputs=[images, concept_sentence],
+        outputs=output_components
+    )
+    images.delete(
+        load_captioning,
+        inputs=[images, concept_sentence],
+        outputs=output_components
+    )
+    images.clear(
+        hide_captioning,
+        outputs=[captioning_area, start]
+    )
+    max_train_epochs.change(
+        fn=update_total_steps,
+        inputs=[max_train_epochs, num_repeats, images],
+        outputs=[total_steps]
+    )
+    num_repeats.change(
+        fn=update_total_steps,
+        inputs=[max_train_epochs, num_repeats, images],
+        outputs=[total_steps]
+    )
+    images.upload(
+        fn=update_total_steps,
+        inputs=[max_train_epochs, num_repeats, images],
+        outputs=[total_steps]
+    )
+    images.delete(
+        fn=update_total_steps,
+        inputs=[max_train_epochs, num_repeats, images],
+        outputs=[total_steps]
+    )
+    images.clear(
+        fn=update_total_steps,
+        inputs=[max_train_epochs, num_repeats, images],
+        outputs=[total_steps]
+    )
+    concept_sentence.change(fn=update_sample, inputs=[concept_sentence], outputs=sample_prompts)
+    start.click(fn=create_dataset, inputs=[dataset_folder, resolution, images] + caption_list, outputs=dataset_folder).then(
+        fn=start_training,
+        inputs=[
+            base_model,
+            lora_name,
+            train_script,
+            train_config,
+            sample_prompts,
+        ],
+        outputs=terminal,
+    )
+    do_captioning.click(fn=run_captioning, inputs=[images, concept_sentence] + caption_list, outputs=caption_list)
+    demo.load(fn=loaded, js=js, outputs=[hf_token, hf_login, hf_logout, repo_owner])
+    refresh.click(update, inputs=listeners, outputs=[train_script, train_config, dataset_folder])
+if __name__ == "__main__":
+    cwd = os.path.dirname(os.path.abspath(__file__))
+    demo.launch(debug=True, show_error=True, allowed_paths=[cwd])

docker-compose.yml

@@ -0,0 +1,28 @@
+services:
+
+  fluxgym:
+    build:
+      context: .
+      # change the dockerfile to Dockerfile.cuda12.4 if you are running CUDA 12.4 drivers otherwise leave as is
+      dockerfile: Dockerfile
+    image: fluxgym
+    container_name: fluxgym
+    ports:
+      - 7860:7860
+    environment:
+      - PUID=${PUID:-1000}
+      - PGID=${PGID:-1000}
+    volumes:
+      - /etc/localtime:/etc/localtime:ro
+      - /etc/timezone:/etc/timezone:ro
+      - ./:/app/fluxgym
+    stop_signal: SIGKILL
+    tty: true
+    deploy:
+      resources:
+        reservations:
+          devices:
+          - driver: nvidia
+            count: all
+            capabilities: [gpu]
+    restart: unless-stopped

install.js

@@ -0,0 +1,96 @@
+module.exports = {
+  run: [
+    {
+      method: "shell.run",
+      params: {
+        venv: "env",
+        message: [
+          "git config --global --add safe.directory '*'",
+          "git clone -b sd3 https://github.com/kohya-ss/sd-scripts"
+        ]
+      }
+    },
+    {
+      method: "shell.run",
+      params: {
+        path: "sd-scripts",
+        venv: "../env",
+        message: [
+          "uv pip install -r requirements.txt",
+        ]
+      }
+    },
+    {
+      method: "shell.run",
+      params: {
+        venv: "env",
+        message: [
+          "pip uninstall -y diffusers[torch] torch torchaudio torchvision",
+          "uv pip install -r requirements.txt",
+        ]
+      }
+    },
+    {
+      method: "script.start",
+      params: {
+        uri: "torch.js",
+        params: {
+          venv: "env",
+          // xformers: true   // uncomment this line if your project requires xformers
+        }
+      }
+    },
+    {
+      method: "fs.link",
+      params: {
+        drive: {
+          vae: "models/vae",
+          clip: "models/clip",
+          unet: "models/unet",
+          loras: "outputs",
+        },
+        peers: [
+          "https://github.com/pinokiofactory/stable-diffusion-webui-forge.git",
+          "https://github.com/pinokiofactory/comfy.git",
+          "https://github.com/cocktailpeanutlabs/comfyui.git",
+          "https://github.com/cocktailpeanutlabs/fooocus.git",
+          "https://github.com/cocktailpeanutlabs/automatic1111.git",
+        ]
+      }
+    },
+//    {
+//      method: "fs.download",
+//      params: {
+//        uri: [
+//          "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/clip_l.safetensors?download=true",
+//          "https://huggingface.co/comfyanonymous/flux_text_encoders/resolve/main/t5xxl_fp16.safetensors?download=true",
+//        ],
+//        dir: "models/clip"
+//      }
+//    },
+//    {
+//      method: "fs.download",
+//      params: {
+//        uri: [
+//          "https://huggingface.co/cocktailpeanut/xulf-dev/resolve/main/ae.sft?download=true",
+//        ],
+//        dir: "models/vae"
+//      }
+//    },
+//    {
+//      method: "fs.download",
+//      params: {
+//        uri: [
+//          "https://huggingface.co/cocktailpeanut/xulf-dev/resolve/main/flux1-dev.sft?download=true",
+//        ],
+//        dir: "models/unet"
+//      }
+//    },
+    {
+      method: "fs.link",
+      params: {
+        venv: "env"
+      }
+    }
+  ]
+}

models.yaml

@@ -0,0 +1,27 @@
+# Add your own model here
+# <name that will show up on the dropdown>:
+#   repo: <the huggingface repo ID to pull from>
+#   base: <the model used to run inference with (The Huggingface "Inference API" widget will use this to generate demo images)>
+#   license: <follow the other examples. Any model inherited from DEV should use the dev license, schenll is apache-2.0>
+#   license_name: <follow the other examples. only needed for dev inherited models>
+#   license_link: <follow the other examples. only needed for dev inherited models>
+#   file: <the file name within the huggingface repo>
+flux-dev:
+    repo: cocktailpeanut/xulf-dev
+    base: black-forest-labs/FLUX.1-dev
+    license: other
+    license_name: flux-1-dev-non-commercial-license
+    license_link: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md
+    file: flux1-dev.sft
+flux-schnell:
+    repo: black-forest-labs/FLUX.1-schnell
+    base: black-forest-labs/FLUX.1-schnell
+    license: apache-2.0
+    file: flux1-schnell.safetensors
+bdsqlsz/flux1-dev2pro-single:
+    repo: bdsqlsz/flux1-dev2pro-single
+    base: black-forest-labs/FLUX.1-dev
+    license: other
+    license_name: flux-1-dev-non-commercial-license
+    license_link: https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md
+    file: flux1-dev2pro.safetensors

pinokio.js

@@ -0,0 +1,95 @@
+const path = require('path')
+module.exports = {
+  version: "3.2",
+  title: "fluxgym",
+  description: "[NVIDIA Only] Dead simple web UI for training FLUX LoRA with LOW VRAM support (From 12GB)",
+  icon: "icon.png",
+  menu: async (kernel, info) => {
+    let installed = info.exists("env")
+    let running = {
+      install: info.running("install.js"),
+      start: info.running("start.js"),
+      update: info.running("update.js"),
+      reset: info.running("reset.js")
+    }
+    if (running.install) {
+      return [{
+        default: true,
+        icon: "fa-solid fa-plug",
+        text: "Installing",
+        href: "install.js",
+      }]
+    } else if (installed) {
+      if (running.start) {
+        let local = info.local("start.js")
+        if (local && local.url) {
+          return [{
+            default: true,
+            icon: "fa-solid fa-rocket",
+            text: "Open Web UI",
+            href: local.url,
+          }, {
+            icon: 'fa-solid fa-terminal',
+            text: "Terminal",
+            href: "start.js",
+          }, {
+            icon: "fa-solid fa-flask",
+            text: "Outputs",
+            href: "outputs?fs"
+          }]
+        } else {
+          return [{
+            default: true,
+            icon: 'fa-solid fa-terminal',
+            text: "Terminal",
+            href: "start.js",
+          }]
+        }
+      } else if (running.update) {
+        return [{
+          default: true,
+          icon: 'fa-solid fa-terminal',
+          text: "Updating",
+          href: "update.js",
+        }]
+      } else if (running.reset) {
+        return [{
+          default: true,
+          icon: 'fa-solid fa-terminal',
+          text: "Resetting",
+          href: "reset.js",
+        }]
+      } else {
+        return [{
+          default: true,
+          icon: "fa-solid fa-power-off",
+          text: "Start",
+          href: "start.js",
+        }, {
+          icon: "fa-solid fa-flask",
+          text: "Outputs",
+          href: "sd-scripts/fluxgym/outputs?fs"
+        }, {
+          icon: "fa-solid fa-plug",
+          text: "Update",
+          href: "update.js",
+        }, {
+          icon: "fa-solid fa-plug",
+          text: "Install",
+          href: "install.js",
+        }, {
+          icon: "fa-regular fa-circle-xmark",
+          text: "Reset",
+          href: "reset.js",
+        }]
+      }
+    } else {
+      return [{
+        default: true,
+        icon: "fa-solid fa-plug",
+        text: "Install",
+        href: "install.js",
+      }]
+    }
+  }
+}

pinokio_meta.json

@@ -0,0 +1,39 @@
+{
+  "posts": [
+    "https://x.com/cocktailpeanut/status/1851721405408166064",
+    "https://x.com/cocktailpeanut/status/1835719701172756592",
+    "https://x.com/LikeToasters/status/1834258975384092858",
+    "https://x.com/cocktailpeanut/status/1834245329627009295",
+    "https://x.com/jkch0205/status/1834003420132614450",
+    "https://x.com/huwhitememes/status/1834074992209699132",
+    "https://x.com/GorillaRogueGam/status/1834148656791888139",
+    "https://x.com/cocktailpeanut/status/1833964839519068303",
+    "https://x.com/cocktailpeanut/status/1833935061907079521",
+    "https://x.com/cocktailpeanut/status/1833940728881242135",
+    "https://x.com/cocktailpeanut/status/1833881392482066638",
+    "https://x.com/Alone1Moon/status/1833348850662445369",
+    "https://x.com/_f_ai_9/status/1833485349995397167",
+    "https://x.com/intocryptoast/status/1833061082862412186",
+    "https://x.com/cocktailpeanut/status/1833888423716827321",
+    "https://x.com/cocktailpeanut/status/1833884852992516596",
+    "https://x.com/cocktailpeanut/status/1833885335077417046",
+    "https://x.com/NiwonArt/status/1833565746624139650",
+    "https://x.com/cocktailpeanut/status/1833884361986380117",
+    "https://x.com/NiwonArt/status/1833599399764889685",
+    "https://x.com/LikeToasters/status/1832934391217045913",
+    "https://x.com/cocktailpeanut/status/1832924887456817415",
+    "https://x.com/cocktailpeanut/status/1832927154536902897",
+    "https://x.com/YabaiHamster/status/1832697724690386992",
+    "https://x.com/cocktailpeanut/status/1832747889497366706",
+    "https://x.com/PhotogenicWeekE/status/1832720544959185202",
+    "https://x.com/zuzaritt/status/1832748542164652390",
+    "https://x.com/foxyy4i/status/1832764883710185880",
+    "https://x.com/waynedahlberg/status/1832226132999213095",
+    "https://x.com/PhotoGarrido/status/1832214644515041770",
+    "https://x.com/cocktailpeanut/status/1832787205774786710",
+    "https://x.com/cocktailpeanut/status/1832151307198541961",
+    "https://x.com/cocktailpeanut/status/1832145996014612735",
+    "https://x.com/cocktailpeanut/status/1832084951115972653",
+    "https://x.com/cocktailpeanut/status/1832091112086843684"
+  ]
+}

requirements.txt

@@ -0,0 +1,35 @@
+safetensors
+git+https://github.com/huggingface/diffusers.git
+gradio_logsview@https://huggingface.co/spaces/cocktailpeanut/gradio_logsview/resolve/main/gradio_logsview-0.0.17-py3-none-any.whl
+transformers
+lycoris-lora==1.8.3
+flatten_json
+pyyaml
+oyaml
+tensorboard
+kornia
+invisible-watermark
+einops
+accelerate
+toml
+albumentations
+pydantic
+omegaconf
+k-diffusion
+open_clip_torch
+timm
+prodigyopt
+controlnet_aux==0.0.7
+python-dotenv
+bitsandbytes
+hf_transfer
+lpips
+pytorch_fid
+optimum-quanto
+sentencepiece
+huggingface_hub
+peft
+gradio
+python-slugify
+imagesize
+pydantic==2.9.2

reset.js

@@ -0,0 +1,13 @@
+module.exports = {
+  run: [{
+    method: "fs.rm",
+    params: {
+      path: "sd-scripts"
+    }
+  }, {
+    method: "fs.rm",
+    params: {
+      path: "env"
+    }
+  }]
+}

start.js

@@ -0,0 +1,37 @@
+module.exports = {
+  daemon: true,
+  run: [
+    {
+      method: "shell.run",
+      params: {
+        venv: "env",                // Edit this to customize the venv folder path
+        env: {
+          LOG_LEVEL: "DEBUG",
+          CUDA_VISIBLE_DEVICES: "0"
+        },                   // Edit this to customize environment variables (see documentation)
+        message: [
+          "python app.py",    // Edit with your custom commands
+        ],
+        on: [{
+          // The regular expression pattern to monitor.
+          // When this pattern occurs in the shell terminal, the shell will return,
+          // and the script will go onto the next step.
+          "event": "/http:\/\/\\S+/",   
+
+          // "done": true will move to the next step while keeping the shell alive.
+          // "kill": true will move to the next step after killing the shell.
+          "done": true
+        }]
+      }
+    },
+    {
+      // This step sets the local variable 'url'.
+      // This local variable will be used in pinokio.js to display the "Open WebUI" tab when the value is set.
+      method: "local.set",
+      params: {
+        // the input.event is the regular expression match object from the previous step
+        url: "{{input.event[0]}}"
+      }
+    }
+  ]
+}

torch.js

@@ -0,0 +1,75 @@
+module.exports = {
+  run: [
+    // windows nvidia
+    {
+      "when": "{{platform === 'win32' && gpu === 'nvidia'}}",
+      "method": "shell.run",
+      "params": {
+        "venv": "{{args && args.venv ? args.venv : null}}",
+        "path": "{{args && args.path ? args.path : '.'}}",
+        "message": "uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 --force-reinstall"
+
+      }
+    },
+    // windows amd
+    {
+      "when": "{{platform === 'win32' && gpu === 'amd'}}",
+      "method": "shell.run",
+      "params": {
+        "venv": "{{args && args.venv ? args.venv : null}}",
+        "path": "{{args && args.path ? args.path : '.'}}",
+        "message": "uv pip install torch-directml torchaudio torchvision"
+      }
+    },
+    // windows cpu
+    {
+      "when": "{{platform === 'win32' && (gpu !== 'nvidia' && gpu !== 'amd')}}",
+      "method": "shell.run",
+      "params": {
+        "venv": "{{args && args.venv ? args.venv : null}}",
+        "path": "{{args && args.path ? args.path : '.'}}",
+        "message": "uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --force-reinstall"
+      }
+    },
+    // mac
+    {
+      "when": "{{platform === 'darwin'}}",
+      "method": "shell.run",
+      "params": {
+        "venv": "{{args && args.venv ? args.venv : null}}",
+        "path": "{{args && args.path ? args.path : '.'}}",
+        "message": "uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --force-reinstall"
+      }
+    },
+    // linux nvidia
+    {
+      "when": "{{platform === 'linux' && gpu === 'nvidia'}}",
+      "method": "shell.run",
+      "params": {
+        "venv": "{{args && args.venv ? args.venv : null}}",
+        "path": "{{args && args.path ? args.path : '.'}}",
+        "message": "uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 --force-reinstall"
+      }
+    },
+    // linux rocm (amd)
+    {
+      "when": "{{platform === 'linux' && gpu === 'amd'}}",
+      "method": "shell.run",
+      "params": {
+        "venv": "{{args && args.venv ? args.venv : null}}",
+        "path": "{{args && args.path ? args.path : '.'}}",
+        "message": "uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1 --force-reinstall"
+      }
+    },
+    // linux cpu
+    {
+      "when": "{{platform === 'linux' && (gpu !== 'amd' && gpu !=='nvidia')}}",
+      "method": "shell.run",
+      "params": {
+        "venv": "{{args && args.venv ? args.venv : null}}",
+        "path": "{{args && args.path ? args.path : '.'}}",
+        "message": "uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu --force-reinstall"
+      }
+    }
+  ]
+}

update.js

@@ -0,0 +1,46 @@
+module.exports = {
+  run: [{
+    method: "shell.run",
+    params: {
+      message: "git pull"
+    }
+  }, {
+    method: "shell.run",
+    params: {
+      path: "sd-scripts",
+      message: "git pull"
+    }
+  }, {
+    method: "shell.run",
+    params: {
+      path: "sd-scripts",
+      venv: "../env",
+      message: [
+        "uv pip install -r requirements.txt",
+      ]
+    }
+  }, {
+    method: "shell.run",
+    params: {
+      venv: "env",
+      message: [
+        "pip uninstall -y diffusers[torch] torch torchaudio torchvision",
+        "uv pip install -r requirements.txt",
+      ]
+    }
+  }, {
+    method: "script.start",
+    params: {
+      uri: "torch.js",
+      params: {
+        venv: "env",
+        // xformers: true   // uncomment this line if your project requires xformers
+      }
+    }
+  }, {
+    method: "fs.link",
+    params: {
+      venv: "env"
+    }
+  }]
+}