AI-ով աշխատող Image Generation API ծառայություն FLUX-ով, Python-ով և Diffusers-ով. Արագ ուղեցույց

Այս հոդվածում մենք ձեզ կառաջարկենք ստեղծել ձեր սեփական 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 նշանաբանը գործարկելու համար
• 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-ի պաշտոնական ուղեցույցին :

Դուք նաև պետք է համոզվեք, որ ունեք առնվազն 12 ԳԲ VRAM, եթե պլանավորում եք գործարկել FLUX-ը GPU սարքի վրա: Կամ առնվազն 12 ԳԲ օպերատիվ հիշողություն՝ պրոցեսորի/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 GPU-ների համար) կամ mps (Apple-ի 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 մոդելը՝ օգտագործելով diffusers գրադարանը: Մոդելը, որը մենք օգտագործում ենք 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 ԳԲ-ից պակաս օպերատիվ հիշողություն: Ուշադրության կտրումը նվազեցնում է հիշողության օգտագործումը պատկերների ստեղծման ժամանակ, ինչը կարևոր է սահմանափակ ռեսուրսներ ունեցող սարքերի համար:

Քայլ 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-ն հայտնի շրջանակ է Python-ով վեբ API-ներ ստեղծելու համար: Այս դեպքում մենք օգտագործում ենք այն սերվեր ստեղծելու համար, որը կարող է ընդունել պատկերների ստեղծման հարցումներ: Մենք նաև օգտագործում ենք GZip միջնակարգ ծրագիր՝ պատասխանը սեղմելու համար, ինչը հատկապես օգտակար է պատկերները base64 ձևաչափով հետ ուղարկելու ժամանակ:

Արտադրական միջավայրում դուք կարող եք պահպանել ստեղծված պատկերները S3 դույլով կամ այլ ամպային պահեստում և վերադարձնել URL-ները 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)