|
|
| from __future__ import annotations |
|
|
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
| from diffusers import ( |
| StableDiffusionPipeline, |
| UNet2DConditionModel, |
| DPMSolverMultistepScheduler, |
| ) |
|
|
| from setup_models import setup_arc2face |
|
|
| BASE_MODEL = "stable-diffusion-v1-5/stable-diffusion-v1-5" |
|
|
|
|
| def load_arc2face(device="cuda"): |
|
|
| |
| repo_path, model_dir = setup_arc2face() |
|
|
| repo_path = Path(repo_path).resolve() |
| model_dir = Path(model_dir).resolve() |
|
|
| if str(repo_path) not in sys.path: |
| sys.path.insert(0, str(repo_path)) |
|
|
| from arc2face import ( |
| CLIPTextModelWrapper, |
| project_face_embs, |
| ) |
|
|
| device = torch.device(device) |
| dtype = ( |
| torch.float16 |
| if device.type == "cuda" |
| else torch.float32 |
| ) |
|
|
| |
| encoder = CLIPTextModelWrapper.from_pretrained( |
| str(model_dir), |
| subfolder="encoder", |
| torch_dtype=dtype, |
| ) |
|
|
| |
| unet = UNet2DConditionModel.from_pretrained( |
| str(model_dir), |
| subfolder="arc2face", |
| torch_dtype=dtype, |
| ) |
|
|
| |
| pipe = StableDiffusionPipeline.from_pretrained( |
| BASE_MODEL, |
| text_encoder=encoder, |
| unet=unet, |
| torch_dtype=dtype, |
| safety_checker=None, |
| ) |
|
|
| pipe.scheduler = DPMSolverMultistepScheduler.from_config( |
| pipe.scheduler.config |
| ) |
|
|
| pipe = pipe.to(device) |
| pipe.set_progress_bar_config(disable=True) |
|
|
| return pipe, dtype |
|
|
|
|
| @torch.inference_mode() |
| def generate_image( |
| pipe, |
| identity_embedding, |
| seed: int, |
| dtype, |
| ): |
|
|
| identity_embedding = F.normalize( |
| identity_embedding.to( |
| device=pipe.device, |
| dtype=dtype, |
| ), |
| p=2, |
| dim=-1, |
| ) |
|
|
| from arc2face import project_face_embs |
|
|
| |
| |
| prompt_embeds = project_face_embs( |
| pipe, |
| identity_embedding, |
| ) |
|
|
| generator = torch.Generator( |
| device=pipe.device |
| ).manual_seed(int(seed)) |
|
|
| result = pipe( |
| prompt_embeds=prompt_embeds, |
| num_inference_steps=25, |
| guidance_scale=3.0, |
| generator=generator, |
| ) |
|
|
| return result.images[0].convert("RGB") |
|
|
