briaai / nfl-body-hd-canny

About

this is the modle currently (12/05/2024) used in our production pipeline. TODO: ASK TAL

Training process

We started from the briaai/BRIA-2.2-HD model (starting from checkpoint 1223000) and finetune it. TODO: ASK TAL

Dataset

TODO: ASK TAL The dataset that was use to train this modle was is located in s3://bria-lora/ds/tal_pair_data/data_pairs/data_pairs_HD_filtered/ the images are at size (1088*2048) and with grey bg color (128,128,128) the caption for the images was: "An illustration of the NFL player: \"{player_name}\" from the team: \"{team_name}\"."

Training parameters

the training script can be found on spring/Infra/project_x/stable_xl/accelerate_launch_prod_xl_us_west.py

lr: 2e-05
batch: 4
accumulate: 2

Evaluation

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    StableDiffusionXLControlNetImg2ImgPipeline,
    EulerAncestralDiscreteScheduler
)
import torch
from PIL import Image

def canny_image(original_image: Image.Image) -> Image.Image:
    canny_image = cv2.Canny(np.array(original_image), 50, 100)
    canny_image = canny_image[:, :, None]
    canny_image = np.concatenate([canny_image, canny_image, canny_image], axis=2)
    return Image.fromarray(canny_image)

scheduler = EulerAncestralDiscreteScheduler(
            beta_start=0.00085,
            beta_end=0.012,
            beta_schedule="scaled_linear",
            num_train_timesteps=1000,
            steps_offset=1,
        )

controlnet = ControlNetModel.from_pretrained("briaai/nfl-body-hd-canny", torch_dtype=torch.float16)
vae = AutoencoderKL.from_pretrained(
            "madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16
        )
pipe = StableDiffusionXLControlNetImg2ImgPipeline.from_pretrained(
            "briaai/nfl_HD",
            torch_dtype=torch.float16,
            controlnet=controlnet,
            scheduler=scheduler,
            vae = vae
        )
pipe.to("cuda")
player_name = "Gal Davidi"
team_name = "Indianapolis Colts"
prompt = "An illustration of the NFL player: \"{player_name}\" from the team: \"{team_name}\"."
negative_prompt = "Logo,Watermark,Text,Ugly,Morbid,Extra fingers,Poorly drawn hands,Mutation,Blurry,Extra limbs,Gross proportions,Missing arms,Mutated hands,Long neck,Duplicate,Mutilated,Mutilated hands,Poorly drawn face,Deformed,Bad anatomy,Cloned face,Malformed limbs,Missing legs,Too many fingers"
img = Image.open("<original image>")
control_image = canny_image(img)
images = pipe(prompt=prompt, negative_prompt=negative_prompt, height=2048, width=1088, image=img, control_image=img,).images[0]