INNER CODE UNIT · Python

main

Shilin-LU/MACE · data_preparation_transformers.py:9

def main(conf):
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # generate 8 images per concept using the original model for performing erasure
    if conf.MACE.generate_data:
        inference(OmegaConf.create({
            "pretrained_model_name_or_path": 'CompVis/stable-diffusion-v1-4',
            "multi_concept": conf.MACE.multi_concept,
            "generate_training_data": True,
            "device": device,
            "steps": 30,
            "output_dir": conf.MACE.input_data_dir,
        }))

    # get and save masks for each image
    if conf.MACE.use_gsam_mask:
        detector_id = "IDEA-Research/grounding-dino-base"

View source record →