چرا Stable Diffusion با JAX/Flax؟

JAX یک کتابخانه‌ی محاسبات عددی است که به‌صورت موازی روی پردازنده‌های GPU و TPU اجرا می‌شود و برای مدل‌های یادگیری عمیق سرعت بالایی ارائه می‌دهد. Flax به‌عنوان بردر فرمت باز برای ساخت ساختارهای مدل در JAX بهینه‌سازی را امکان‌پذیر می‌کند.

  • سرعت بالاتر نسبت به PyTorch در بعضی شرایط.
  • پشتیبانی از TPU و دسترسی آسان در Google Colab.
  • سازگاری کامل با کتابخانه Hugging Face Diffusers نسخه 0.5.1 و بالاتر.
  • قابلیت مقیاس‌پذیری افقی برای پردازش تصاویر بزرگ.

نصب و راه‌اندازی مدل روی TPU

برای استفاده از TPU در گُوگل کولاب کافی است از منوی Runtime گزینه‌ی “Change runtime type” را انتخاب کرده و تحت Hardware accelerator را به TPU تنظیم کنید. سپس کتابخانه‌ی diffusers را با دستور زیر نصب کنید:
pip install diffusers[torch] --upgrade

بعد از نصب، با ایجاد یک شیکلاس DiffusionPipeline را با پارامترهای Flax برمی‌گردانید و سریعاً امکان تولید تصاویر را می‌یابید:

تجزیه و تحلیل عملکرد و نکات اجرایی

مدلی که با استفاده از JAX/Flax ساخته می‌شود، در خط فرمان‌های متوالی سرعت 2-3 برابر سریع‌تر از نسخه‌ی PyTorch پیچیده است. نکات زیر را در ذهن داشته باشید:

  • مقیاس‌پذیری بار ورودی (batch size) را به صورت دنباله‌های کوچک تغییر دهید.
  • محدود کردن مصرف حافظه با پارامتر «dtype = float16» درون مسیر کار کمک‌کننده است.
  • برای دانلود مدل‌های بزرگ از حافظه‌ی درون-کنسروشن (in‑memory storage) در TPU استفاده کنید.

برای دریافت بیشتر مطالب و مستندات به پست اصلی Hugging Face مراجعه فرمایید.

پرسش‌های پرتکرار

کدام نسخه Diffusers برای اجرای Stable Diffusion با JAX/Flax ضروری است؟

از نسخه 0.5.1 به بالا، Diffusers به‌طور پیش‌فرض پشتیبانی از Flax را فراهم می‌کند.

می‌توانم در قالب Colab، مستقیماً از TPU برای تولید تصویر استفاده کنم؟

بله، کافی است در بخش Runtime گزینه TPU را فعال کنید، کتابخانه‌های مورد نیاز را نصب کنید و مدل را با JAX/Flax فراخوانی کنید.