mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-02-01 17:14:58 -05:00
b3d60bd56a
Similar to the existing node, but without any resizing and with a revised model loading API that uses the model manager. All code related to the invocation now lives in the Invoke repo. Unfortunately, this includes a whole git repo for EfficientNet. I believe we could use the package `timm` instead of this, but it's beyond me.
22 lines
597 B
Python
22 lines
597 B
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
from .submodules.encoder import Encoder
|
|
from .submodules.decoder import Decoder
|
|
|
|
|
|
class NNET(nn.Module):
|
|
def __init__(self, args):
|
|
super(NNET, self).__init__()
|
|
self.encoder = Encoder()
|
|
self.decoder = Decoder(args)
|
|
|
|
def get_1x_lr_params(self): # lr/10 learning rate
|
|
return self.encoder.parameters()
|
|
|
|
def get_10x_lr_params(self): # lr learning rate
|
|
return self.decoder.parameters()
|
|
|
|
def forward(self, img, **kwargs):
|
|
return self.decoder(self.encoder(img), **kwargs) |