שירות API ליצירת תמונה מבוססת בינה מלאכותית עם FLUX, Python ומפזרים: מדריך מהיר

במאמר זה, נדריך אותך דרך יצירת שרת FLUX משלך באמצעות Python. שרת זה יאפשר לך ליצור תמונות על סמך הנחיות טקסט באמצעות API פשוט. בין אם אתה מפעיל שרת זה לשימוש אישי או פורס אותו כחלק מיישום ייצור, מדריך זה יעזור לך להתחיל.

FLUX (של Black Forest Labs ) כבשה בסערה את עולם יצירת תמונות AI בחודשים האחרונים. לא רק שהיא ניצחה את Stable Diffusion (מלך הקוד הפתוח הקודם) במדדים רבים, אלא שהיא גם עלתה על דגמים קנייניים כמו Dall-E או Midjourney בכמה מדדים.

אבל איך היית ממשיך להשתמש ב-FLUX באחת מהאפליקציות שלך? אפשר לחשוב על שימוש במארחים חסרי שרת כמו Replicate ואחרים, אבל אלה עלולים להתייקר מהר מאוד, ואולי לא מספקים את הגמישות הדרושה לך. זה המקום שבו יצירת שרת FLUX מותאם אישית משלך שימושית.

דרישות מוקדמות

לפני שצולל לתוך הקוד, בואו נוודא שיש לכם את הכלים והספריות הדרושים:

• Python: יהיה עליך להתקין את Python 3 במחשב שלך, רצוי גרסה 3.10.
• torch : מסגרת הלמידה העמוקה שבה נשתמש להפעלת FLUX.
• diffusers : מספק גישה לדגם FLUX.
• transformers : תלות נדרשת של מפזרים.
• sentencepiece : נדרש להפעלת ה-FLUX tokenizer
• protobuf : נדרש להפעלת FLUX
• accelerate : עוזר לטעון את דגם FLUX בצורה יעילה יותר במקרים מסוימים.
• fastapi : מסגרת ליצירת שרת אינטרנט שיכול לקבל בקשות ליצירת תמונות.
• uvicorn : נדרש להפעלת שרת FastAPI.
• psutil : מאפשר לנו לבדוק כמה זיכרון RAM יש במחשב שלנו.

אתה יכול להתקין את כל הספריות על ידי הפעלת הפקודה הבאה: pip install torch diffusers transformers sentencepiece protobuf accelerate fastapi uvicorn .

אם אתה משתמש ב-Mac עם שבב M1 או M2, עליך להגדיר PyTorch עם Metal לביצועים מיטביים. עקוב אחר המדריך הרשמי של PyTorch with Metal לפני שתמשיך.

תצטרך גם לוודא שיש לך לפחות 12 GB של VRAM אם אתה מתכנן להפעיל את FLUX על מכשיר GPU. או לפחות 12 GB של זיכרון RAM להפעלה על CPU/MPS (שיהיה איטי יותר).

שלב 1: הגדרת הסביבה

הבה נתחיל את הסקריפט על ידי בחירת המכשיר המתאים להפעיל הסקה על סמך החומרה שבה אנו משתמשים.

 device = 'cuda' # can also be 'cpu' or 'mps' import os # MPS support in PyTorch is not yet fully implemented if device == 'mps': os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" import torch if device == 'mps' and not torch.backends.mps.is_available(): raise Exception("Device set to MPS, but MPS is not available") elif device == 'cuda' and not torch.cuda.is_available(): raise Exception("Device set to CUDA, but CUDA is not available")

אתה יכול לציין cpu , cuda (עבור NVIDIA GPUs) או mps (עבור Metal Performance Shaders של אפל). לאחר מכן, הסקריפט בודק אם המכשיר שנבחר זמין ומעלה חריג אם לא.

שלב 2: טעינת דגם FLUX

לאחר מכן, נטען את דגם FLUX. נטען את הדגם ב-fp16 דיוק שיחסוך לנו קצת זיכרון בלי הפסד גדול באיכות.

בשלב זה, ייתכן שתתבקשו לבצע אימות באמצעות HuggingFace, מכיוון שדגם ה-FLUX הוא סגור. על מנת לאמת בהצלחה, תצטרך ליצור חשבון HuggingFace, לעבור לדף הדגם, לקבל את התנאים, ולאחר מכן ליצור אסימון HuggingFace מהגדרות החשבון שלך ולהוסיף אותו במחשב שלך בתור משתנה הסביבה HF_TOKEN .

 from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline import psutil model_name = "black-forest-labs/FLUX.1-dev" print(f"Loading {model_name} on {device}") pipeline = FluxPipeline.from_pretrained( model_name, # Diffusion models are generally trained on fp32, but fp16 # gets us 99% there in terms of quality, with just half the (V)RAM torch_dtype=torch.float16, # Ensure we don't load any dangerous binary code use_safetensors=True # We are using Euler here, but you can also use other samplers scheduler=FlowMatchEulerDiscreteScheduler() ).to(device)

כאן, אנו טוענים את מודל FLUX באמצעות ספריית המפזרים. הדגם שבו אנו משתמשים הוא black-forest-labs/FLUX.1-dev , נטען ב-fp16 דיוק.

קיים גם דגם מזוקק בשלבי זמן בשם FLUX Schnell בעל הסקה מהירה יותר, אך מפיק תמונות פחות מפורטות, וכן דגם FLUX Pro שהוא במקור סגור. נשתמש כאן בלוח הזמנים של אוילר, אבל אתה יכול להתנסות בזה. אתה יכול לקרוא עוד על מתזמנים כאן . מכיוון שיצירת תמונות עשויה להיות עתירת משאבים, חיוני לייעל את השימוש בזיכרון, במיוחד כאשר פועל על מעבד או מכשיר עם זיכרון מוגבל.

 # Recommended if running on MPS or CPU with < 64 GB of RAM total_memory = psutil.virtual_memory().total total_memory_gb = total_memory / (1024 ** 3) if (device == 'cpu' or device == 'mps') and total_memory_gb < 64: print("Enabling attention slicing") pipeline.enable_attention_slicing()

קוד זה בודק את סך הזיכרון הזמין ומאפשר חיתוך תשומת לב אם למערכת יש פחות מ-64 GB של זיכרון RAM. חיתוך תשומת הלב מפחית את השימוש בזיכרון במהלך יצירת תמונה, דבר חיוני עבור מכשירים עם משאבים מוגבלים.

שלב 3: יצירת ה-API עם FastAPI

לאחר מכן, נגדיר את שרת FastAPI, שיספק API ליצירת תמונות.

 from fastapi import FastAPI, HTTPException from pydantic import BaseModel, Field, conint, confloat from fastapi.middleware.gzip import GZipMiddleware from io import BytesIO import base64 app = FastAPI() # We will be returning the image as a base64 encoded string # which we will want compressed app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7)

FastAPI היא מסגרת פופולרית לבניית ממשקי API באינטרנט עם Python. במקרה זה, אנו משתמשים בו כדי ליצור שרת שיכול לקבל בקשות ליצירת תמונות. אנחנו גם משתמשים בתוכנת GZip כדי לדחוס את התגובה, וזה שימושי במיוחד כששולחים תמונות בחזרה בפורמט base64.

בסביבת ייצור, ייתכן שתרצה לאחסן את התמונות שנוצרו בדלי S3 או באחסון ענן אחר ולהחזיר את כתובות האתרים במקום המחרוזות המקודדות base64, כדי לנצל את היתרון של CDN ואופטימיזציות אחרות.

שלב 4: הגדרת מודל הבקשה

כעת עלינו להגדיר מודל לבקשות שה-API שלנו יקבל.

 class GenerateRequest(BaseModel): prompt: str seed: conint(ge=0) = Field(..., description="Seed for random number generation") height: conint(gt=0) = Field(..., description="Height of the generated image, must be a positive integer and a multiple of 8") width: conint(gt=0) = Field(..., description="Width of the generated image, must be a positive integer and a multiple of 8") cfg: confloat(gt=0) = Field(..., description="CFG (classifier-free guidance scale), must be a positive integer or 0") steps: conint(ge=0) = Field(..., description="Number of steps") batch_size: conint(gt=0) = Field(..., description="Number of images to generate in a batch")

מודל GenerateRequest זה מגדיר את הפרמטרים הנדרשים ליצירת תמונה. שדה prompt הוא תיאור הטקסט של התמונה שברצונך ליצור. שדות אחרים כוללים את ממדי התמונה, מספר שלבי ההסקה וגודל האצווה.

שלב 5: יצירת נקודת הקצה ליצירת תמונה

כעת, בואו ניצור את נקודת הקצה שתטפל בבקשות ליצירת תמונות.

 @app.post("/") async def generate_image(request: GenerateRequest): # Validate that height and width are multiples of 8 # as required by FLUX if request.height % 8 != 0 or request.width % 8 != 0: raise HTTPException(status_code=400, detail="Height and width must both be multiples of 8") # Always calculate the seed on CPU for deterministic RNG # For a batch of images, seeds will be sequential like n, n+1, n+2, ... generator = [torch.Generator(device="cpu").manual_seed(i) for i in range(request.seed, request.seed + request.batch_size)] images = pipeline( height=request.height, width=request.width, prompt=request.prompt, generator=generator, num_inference_steps=request.steps, guidance_scale=request.cfg, num_images_per_prompt=request.batch_size ).images # Convert images to base64 strings # (for a production app, you might want to store the # images in an S3 bucket and return the URLs instead) base64_images = [] for image in images: buffered = BytesIO() image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") base64_images.append(img_str) return { "images": base64_images, }

נקודת קצה זו מטפלת בתהליך יצירת התמונה. תחילה הוא מאמת שהגובה והרוחב הם כפולות של 8, כנדרש על ידי FLUX. לאחר מכן הוא יוצר תמונות על סמך ההנחיה שסופקה ומחזיר אותן כמחרוזות מקודדות base64.

שלב 6: הפעלת השרת

לבסוף, בואו נוסיף קצת קוד כדי להפעיל את השרת כאשר הסקריפט מופעל.

 @app.on_event("startup") async def startup_event(): print("Image generation server running") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)

קוד זה מתחיל את שרת FastAPI ביציאה 8000, מה שהופך אותו לנגיש לא רק מ- http://localhost:8000 אלא גם ממכשירים אחרים באותה רשת באמצעות כתובת ה-IP של המחשב המארח, הודות לקשירת 0.0.0.0 .

שלב 7: בדיקת השרת שלך באופן מקומי

עכשיו כששרת FLUX שלך פועל, הגיע הזמן לבדוק אותו. אתה יכול להשתמש curl , כלי שורת פקודה לביצוע בקשות HTTP, כדי ליצור אינטראקציה עם השרת שלך:

 curl -X POST "http://localhost:8000/" \ -H "Content-Type: application/json" \ -d '{ "prompt": "A futuristic cityscape at sunset", "seed": 42, "height": 1024, "width": 1024, "cfg": 3.5, "steps": 50, "batch_size": 1 }' | jq -r '.images[0]' | base64 -d > test.png

פקודה זו תעבוד רק על מערכות מבוססות UNIX עם כלי השירות curl , jq ו- base64 מותקנים. ההשלמה עשויה להימשך עד מספר דקות, בהתאם לחומרה המארחת את שרת FLUX.

מַסְקָנָה

מזל טוב! יצרת בהצלחה שרת FLUX משלך באמצעות Python. הגדרה זו מאפשרת לך ליצור תמונות על סמך הנחיות טקסט באמצעות ממשק API פשוט. אם אינך מרוצה מהתוצאות של דגם הבסיס של FLUX, אולי תשקול לכוונן את המודל לביצועים טובים עוד יותר במקרים שימוש ספציפיים .

קוד מלא

אתה עשוי למצוא את הקוד המלא המשמש במדריך זה למטה:

 device = 'cuda' # can also be 'cpu' or 'mps' import os # MPS support in PyTorch is not yet fully implemented if device == 'mps': os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1" import torch if device == 'mps' and not torch.backends.mps.is_available(): raise Exception("Device set to MPS, but MPS is not available") elif device == 'cuda' and not torch.cuda.is_available(): raise Exception("Device set to CUDA, but CUDA is not available") from diffusers import FlowMatchEulerDiscreteScheduler, FluxPipeline import psutil model_name = "black-forest-labs/FLUX.1-dev" print(f"Loading {model_name} on {device}") pipeline = FluxPipeline.from_pretrained( model_name, # Diffusion models are generally trained on fp32, but fp16 # gets us 99% there in terms of quality, with just half the (V)RAM torch_dtype=torch.float16, # Ensure we don't load any dangerous binary code use_safetensors=True, # We are using Euler here, but you can also use other samplers scheduler=FlowMatchEulerDiscreteScheduler() ).to(device) # Recommended if running on MPS or CPU with < 64 GB of RAM total_memory = psutil.virtual_memory().total total_memory_gb = total_memory / (1024 ** 3) if (device == 'cpu' or device == 'mps') and total_memory_gb < 64: print("Enabling attention slicing") pipeline.enable_attention_slicing() from fastapi import FastAPI, HTTPException from pydantic import BaseModel, Field, conint, confloat from fastapi.middleware.gzip import GZipMiddleware from io import BytesIO import base64 app = FastAPI() # We will be returning the image as a base64 encoded string # which we will want compressed app.add_middleware(GZipMiddleware, minimum_size=1000, compresslevel=7) class GenerateRequest(BaseModel): prompt: str seed: conint(ge=0) = Field(..., description="Seed for random number generation") height: conint(gt=0) = Field(..., description="Height of the generated image, must be a positive integer and a multiple of 8") width: conint(gt=0) = Field(..., description="Width of the generated image, must be a positive integer and a multiple of 8") cfg: confloat(gt=0) = Field(..., description="CFG (classifier-free guidance scale), must be a positive integer or 0") steps: conint(ge=0) = Field(..., description="Number of steps") batch_size: conint(gt=0) = Field(..., description="Number of images to generate in a batch") @app.post("/") async def generate_image(request: GenerateRequest): # Validate that height and width are multiples of 8 # as required by FLUX if request.height % 8 != 0 or request.width % 8 != 0: raise HTTPException(status_code=400, detail="Height and width must both be multiples of 8") # Always calculate the seed on CPU for deterministic RNG # For a batch of images, seeds will be sequential like n, n+1, n+2, ... generator = [torch.Generator(device="cpu").manual_seed(i) for i in range(request.seed, request.seed + request.batch_size)] images = pipeline( height=request.height, width=request.width, prompt=request.prompt, generator=generator, num_inference_steps=request.steps, guidance_scale=request.cfg, num_images_per_prompt=request.batch_size ).images # Convert images to base64 strings # (for a production app, you might want to store the # images in an S3 bucket and return the URL's instead) base64_images = [] for image in images: buffered = BytesIO() image.save(buffered, format="PNG") img_str = base64.b64encode(buffered.getvalue()).decode("utf-8") base64_images.append(img_str) return { "images": base64_images, } @app.on_event("startup") async def startup_event(): print("Image generation server running") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)