johnowhitaker/sac_midu_mini

# 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)