(args)
| 94 | |
| 95 | |
| 96 | def main(args): |
| 97 | # --- Distributed Setup --- |
| 98 | rank = int(os.environ.get("RANK", "0")) |
| 99 | world_size = int(os.environ.get("WORLD_SIZE", "1")) |
| 100 | local_rank = int(os.environ.get("LOCAL_RANK", "0")) |
| 101 | |
| 102 | setup_distributed(rank, world_size) |
| 103 | device = torch.device(f"cuda:{local_rank}") |
| 104 | torch.cuda.set_device(device) |
| 105 | |
| 106 | # --- Mixed Precision Setup --- |
| 107 | mixed_precision_dtype = None |
| 108 | if args.mixed_precision == "fp16": |
| 109 | mixed_precision_dtype = torch.float16 |
| 110 | elif args.mixed_precision == "bf16": |
| 111 | mixed_precision_dtype = torch.bfloat16 |
| 112 | |
| 113 | enable_amp = mixed_precision_dtype is not None |
| 114 | |
| 115 | if is_main_process(rank): |
| 116 | print(f"Running distributed evaluation with {world_size} GPUs.") |
| 117 | if enable_amp: |
| 118 | print(f"Using mixed precision: {args.mixed_precision}") |
| 119 | os.makedirs(args.output_dir, exist_ok=True) |
| 120 | if args.save_images: |
| 121 | os.makedirs(os.path.join(args.output_dir, "images"), exist_ok=True) |
| 122 | |
| 123 | results_filepath = os.path.join(args.output_dir, "evaluation_results.jsonl") |
| 124 | |
| 125 | # --- Load Model and Pipeline --- |
| 126 | if is_main_process(rank): |
| 127 | print("Loading model and pipeline...") |
| 128 | |
| 129 | if args.model_type == "sd3": |
| 130 | pipeline = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium") |
| 131 | target_modules = [ |
| 132 | "attn.add_k_proj", |
| 133 | "attn.add_q_proj", |
| 134 | "attn.add_v_proj", |
| 135 | "attn.to_add_out", |
| 136 | "attn.to_k", |
| 137 | "attn.to_out.0", |
| 138 | "attn.to_q", |
| 139 | "attn.to_v", |
| 140 | ] |
| 141 | transformer_lora_config = LoraConfig( |
| 142 | r=32, lora_alpha=64, init_lora_weights="gaussian", target_modules=target_modules |
| 143 | ) |
| 144 | else: |
| 145 | raise ValueError(f"Unsupported model type: {args.model_type}") |
| 146 | |
| 147 | torch.backends.cuda.matmul.allow_tf32 = True |
| 148 | torch.backends.cudnn.allow_tf32 = True |
| 149 | |
| 150 | if args.lora_hf_path: |
| 151 | pipeline.transformer = PeftModel.from_pretrained(pipeline.transformer, args.lora_hf_path) |
| 152 | pipeline.transformer = pipeline.transformer.merge_and_unload() |
| 153 | elif args.checkpoint_path: |
no test coverage detected