Skip to content

collect_model

Define collector for video.

CollectModel

Bases: InstancePlugin

Collect 3D Model for publish.

TODO(@sas): Only supports exports for singular models, expand so that sequences of objs/glbs can be properly taken in.

Source code in client/ayon_comfyui/plugins/publish/collect_model.py
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
class CollectModel(pyblish.api.InstancePlugin):
    """Collect 3D Model for publish.

    TODO(@sas): Only supports exports for singular models,
    expand so that sequences of objs/glbs can be properly taken in.
    """

    order = pyblish.api.CollectorOrder + 0.17
    label = "Collect Generated Model"
    hosts: ClassVar[list[str]] = ["comfyui"]
    families: ClassVar[list[str]] = ["model"]

    default_variant = "Main"

    model_exts: ClassVar[list[str]] = [
        ".gltf",
        ".glb",
        ".obj",
        ".fbx",
        ".stl",
        ".spz",
        ".splat",
        ".ply",
        ".ksplat",
    ]

    def process(self, instance: pyblish.api.Instance):
        host: ComfyUIHost = registered_host()
        image_urls = host.stub.get_publish_node_images(
            instance.data, publish_type=PublishType.MODEL3D
        )

        instance.data["anatomyData"] = instance.context.data["anatomyData"]
        staging_dir = get_instance_staging_dir(instance)
        self.log.info("Outputting model to: %s", staging_dir)

        model_link = next(iter(image_urls))
        if model_link is None:
            self.log.warning("Nothing could be collected. (No url returned.)")
            return

        # Download model
        self.log.debug("Downloading model: %s", model_link)
        parse = urlsplit(model_link)
        self.log.debug(parse)
        query = parse_qs(parse.query)
        self.log.debug(query)
        filename = next(iter(query.get("filename")), None)
        if filename is None:
            self.log.warning(
                "Nothing could be collected. (No filename in query.)"
            )
            return
        if (extension := Path(filename).suffix) not in self.model_exts:
            self.log.warning(
                "Nothing could be collected. "
                "(filename has invalid extension for video.)"
            )
            return

        product_name: str = instance.data["productName"]
        destination = os.path.join(staging_dir, product_name, filename)
        model_file = os.path.join(product_name, filename)
        Path(destination).parent.mkdir(parents=True, exist_ok=True)
        urlretrieve(model_link, destination)  # noqa: S310

        instance.context.data["currentFile"] = model_file

        # creating representation
        instance.data["representations"].append(
            {
                "name": extension[1:],
                "ext": extension[1:],
                "files": model_file,
                "stagingDir": staging_dir,
            }
        )