Training Flux Locally on Mac

<p>

For all those struggling to set this up right now.</p>
<p><em><strong><strong>(rearticulated by A.C.T. soon® from a post/repo by Hughescr and the ai-toolkit Flux training script by Ostris)</strong></strong></em></p>
<p>This workflow is not grounded in Diffusers. However, I have not yet encountered a working Diffusers implementation of local Flux training on Mac/mps. If such a workflow/pipeline exists, I would sincerely appreciate it if someone linked me to it (or/and advised me on implementation details). Such as via <a href="mailto:alekseycalvin@gmail.com" rel="nofollow">alekseycalvin@gmail.com</a>, or a comment somewhere... Like, say, one of my Flux LoRA repos here on Huggingface... (By the way, check them out? To improve Flux Schnell, use <a href="https://huggingface.co/AlekseyCalvin/historic_color_schnell">Historic Color Schnell</a>.)</p>
<p>But to the point, the repo to train locally on Mac is <a href="https://github.com/hughescr/ai-toolkit" rel="nofollow">here</a>, as a somewhat modified branch of Ostris' ai-toolkit training script git.</p>
<p><strong>Below sits the link to the</strong> ai-toolkit <strong>repo modified for MacOS:</strong>
<a href="https://github.com/hughescr/ai-toolkit" rel="nofollow">https://github.com/hughescr/ai-toolkit</a></p>
<p>To be clear, I'm not the person behind this branch, but myself finally stumbled upon it whilst seeking far and wide for many hours any extant Flux training solution adapted for MacOS/sillicon. So, if this works for you, then please thank that prodigious wizard Ostris (the developer of ai-toolkit training scripts), along with this Mac-oriented branch's mysterious author: a certain Hughescr. </p>
<p>Credit and solidarity further extends to all who -- in chronic scepticism of seemingly insurmountable limitations -- stubbornly tinker and quest for options, solutions, and possibilities. </p>
<p>In any case, on the basis of another guide post by Hughescr, plus a few notes/details added from myself for clarity, I just put together the below short guide on setting up local Flux training on Macs using ai-toolkit + the Hughescr branch. </p>
<p><em><strong>Take heed though: without further optimization, this is very unlikely to work on Mac systems with low unified memory!</strong></em> </p>
<p><strong><strong>WORKFLOW to TRAIN FLUX On Mac/OSX/mps:</strong></strong></p>
<p><strong>In Terminal, clone <strong><a href="https://github.com/hughescr/ai-toolkit" rel="nofollow">https://github.com/hughescr/ai-toolkit</a></strong>, following the Linux instructions in the</strong> README <strong>there.</strong></p>
<p><strong>As in:</strong></p>
<pre><code class="language-python">git clone https://github.com/hughescr/ai-toolkit
</code></pre>
<p><strong>Then travel over to the cloned directory:</strong></p>
<pre><code class="language-python">cd ai-toolkit
</code></pre>
<p><strong>Do this:</strong></p>
<pre><code class="language-python">git submodule update --init --recursive
</code></pre>
<p><strong>Then make a virtual environment from the same folder:</strong></p>
<pre><code class="language-python">python3 -m venv venv
</code></pre>
<p><strong>Activate it:</strong></p>
<pre><code class="language-python">source venv/<span class="hljs-built_in">bin</span>/activate
</code></pre>
<p><strong>Install</strong> PyTorch <strong>:</strong></p>
<pre><code class="language-python">pip3 install torch
</code></pre>
<p><strong>Install requirements for the</strong> ai-toolkit, <strong>which should also extend it with certain submodules updated/introduced by Hughescr:</strong></p>
<pre><code class="language-python">pip3 install -r requirements.txt
</code></pre>
<p>**<strong>Here's a list of all that Hughescr introduced to that branch of Ostris'</strong> ai-toolkit training script, <strong>in order to adapt it for Mac OSX (I quote the below from their post):</strong></p>
<p><strong>-- Using</strong> torch.amp <strong>instead of</strong> torch.cuda.amp <strong>(which should work for CUDA too, but will allow MPS to work, using an MPS-compatible GradScaler).</strong></p>
<p><strong>-- Force-using</strong> spawn <strong>instead of</strong> fork <strong>for multiprocessing.</strong></p>
<p><strong>-- Turning off the T5 quantizer, because this won't work on MPS.</strong></p>
<p><strong>-- Forcing the dataloader to have</strong> num_workers=0 <strong>(otherwise the process breaks on Mac).</strong>
<em>This may be done by adding</em> "num_workers=0" <em>to your <strong>config file</strong> for a prospective training: in this context, this would be your variant of one of the</em> (.yaml) <em>template configs from</em>  /ai-toolkit/config/examples. <em>The Hughescr branch of ai-toolkit is supposed to already pre-enforce this particular option, even irrespectively of the config, but it might be better to make doubly sure manually.</em></p>
<p><em>On a side note for those aspiring Flux remodelers who are new to training scripts or relatively code-fresh: The template config file, typically in a file format of either</em> .yaml <em>or</em> .json <em>(such as for</em> <a href="https://github.com/kohya-ss/sd-scripts" rel="nofollow">Kohya ss</a><em>) is an essential component of launching a local training (at least without some GUI interface container/app), and typically carries strict internal formatting rules, corresponding to its data type broadly and/or family/architecture of trainers more specifically. As such, whilst specifying</em> "num_workers=0" <em>or or filling in your training parameters or modifying anything else within a config .yaml (or .json, etc), make sure to match closely the format and syntax found throughout the config template! Else get condemned to exasperating backtracking come runtime.</em> </p>
<p><em>Relatedly, the</em> /ai-toolkit <em>local trainer scripts folder contains a wide range of template configs, not just for training Flux, but for many other sorts of models as well. There's much there to explore and potentially try. Instrumental to our given case, however, are the specific template config</em> .yaml <em>files for</em> <a href="https://huggingface.co/black-forest-labs/FLUX.1-dev">Flux Dev</a> <em>and</em> <a href="/blog/AlekseyCalvin/black-forest-labs/FLUX.1-schnell">Flux Schnell</a>. <em>These configs,</em> train_lora_flux_24gb.yaml <em>and</em> train_lora_flux_schnell_24gb.yaml, <em>are found in the</em> <strong>/config/examples/</strong> <em>subfolder of the cloned-in</em> <strong>/ai-toolkit</strong> <em>folder: the relevant config (for either training Dev or Schnell) is meant to get duplicated by you and thereafter modified for your use with the training script. These configs may then be brought in as an argument to the</em> run.py <em>launcher, if you want to launch the trainer all at once directly from the config and through the Terminal. Or the config could be dragged into/opened via a dedicated UI. The built-in ai-toolkit UI may be launched from the same</em> <strong>/ai-toolkit</strong> <em>folder, via the</em>  flux_train_ui.py  <em>Python executable.)</em></p>
<pre><code>   🌕🌖🌗🌘🌑🌒🌓🌔🌕🌖🌗🌘🌑🌒🌓🌔🌕🌖🌗🌘🌑🌒🌓🌔🌕
</code></pre>
<p><strong>NOW, to run/modify the script, follow further usage instructions here:</strong></p>
<p><a href="https://github.com/hughescr/ai-toolkit#tutorial" rel="nofollow">https://github.com/hughescr/ai-toolkit#tutorial</a></p>
<p><strong>Finally, in order to side-step functions un-implemented in MPS (thus far), one needs to launch the Python executable for training script with the following argument:</strong></p>
<pre><code class="language-python">PYTORCH_ENABLE_MPS_FALLBACK=<span class="hljs-number">1</span>
</code></pre>
<p><strong>This is basically a way of enabling selective/temporary CPU-offload of operations unable to work with MPS/sillicon.</strong></p>
<p><strong>As in:</strong></p>
<pre><code class="language-python">PYTORCH_ENABLE_MPS_FALLBACK=<span class="hljs-number">1</span> python run.py config/your-custom-config-file.yaml
</code></pre>
<p><strong>This should launch the custom training config, and thereby the training itself!</strong></p>
<p><strong>Lastly, just for clarity, and in case anyone reading this is new to manually-launched training, I will reiterate:</strong></p>
<p><strong>To specify stuff like, say,</strong> <em>dataset folder location input, model output folder location, trigger phrase/token, learning rate, optimizer, etc..</em>. <strong>one must duplicate and modify the</strong> .yaml <strong>config file from the</strong> /config <strong>or the</strong>  /config/examples/ <strong>subfolder of your</strong> /ai-toolkit <strong>folder...</strong></p>
<p><strong>NOW GO AND TRY THIS!</strong></p>
<p><em><strong>Sincerely,</strong></em> </p>
<p><em><strong>A.C.T. SOON®</strong></em></p>
