| ```python | |
| # Download the model weights | |
| from huggingface_hub import hf_hub_download | |
| model_path = hf_hub_download(repo_id="johnowhitaker/sac_midu_mini", | |
| filename="midu_model_aesthetic_classifier.pt") | |
| # Load the aesthetic classifier | |
| m = nn.Sequential( | |
| nn.Conv2d(1280, 256, kernel_size=3, padding=1), nn.ReLU(), | |
| nn.MaxPool2d(2, 2), | |
| nn.Conv2d(256, 128, kernel_size=3, padding=1), nn.ReLU(), | |
| nn.AdaptiveAvgPool2d(output_size=(2, 2)), nn.Flatten(), | |
| nn.Linear(128*4, 64), nn.ReLU(), nn.Linear(64, 10)).to(device) | |
| m.load_state_dict(torch.load(model_path)); | |
| # Load the SD pipeline and add a hook | |
| pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base").to(device) | |
| pipe.scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000) | |
| pipe.scheduler.set_timesteps(30) | |
| def hook_fn(module, input, output): | |
| module.output = output | |
| pipe.unet.mid_block.register_forward_hook(hook_fn); | |
| # Now after calling the forward pass of the UNET, you can do | |
| preds = m(pipe.unet.mid_block.output) | |
| ``` |