INNER CODE UNIT · Python

main

Shilin-LU/MACE · inference.py:8

def main(args):

    model_id = args.pretrained_model_name_or_path
    pipe = StableDiffusionPipeline.from_pretrained(model_id).to(args.device)
    pipe.safety_checker = None
    pipe.requires_safety_checker = False
    torch.Generator(device=args.device).manual_seed(42)
    
    if args.generate_training_data:
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
        num_images = 8
        count = 0
        for single_concept in args.multi_concept:
            for c, t in single_concept:
                count += 1
                print(f"Generating training data for concept {count}: {c}...")
                c = c.replace('-', ' ')
                output_folder = f"{args.output_dir}/{c}"

View source record →