File size: 1,067 Bytes
c19ca42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
import navi
from nodes.base_output import BaseOutput, OutputKind

from ...impl.ncnn.model import NcnnModelWrapper
from ...utils.format import format_channel_numbers


class NcnnModelOutput(BaseOutput):
    def __init__(
        self,
        model_type: navi.ExpressionJson = "NcnnNetwork",
        label: str = "Model",
        kind: OutputKind = "generic",
    ):
        super().__init__(model_type, label, kind=kind, associated_type=NcnnModelWrapper)

    def get_broadcast_data(self, value: NcnnModelWrapper):
        return {
            "tags": [
                format_channel_numbers(value.in_nc, value.out_nc),
                f"{value.nf}nf",
                value.fp,
            ]
        }

    def get_broadcast_type(self, value: NcnnModelWrapper):
        return navi.named(
            "NcnnNetwork",
            {
                "scale": value.scale,
                "inputChannels": value.in_nc,
                "outputChannels": value.out_nc,
                "nf": value.nf,
                "fp": navi.literal(value.fp),
            },
        )