Skip to content

collect_image

CollectImage

Bases: InstancePlugin

Collect for publish.

Source code in client/ayon_comfyui/plugins/publish/collect_image.py 
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
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
class CollectImage(pyblish.api.InstancePlugin):
    """Collect  for publish."""

    order = pyblish.api.CollectorOrder + 0.15
    label = "Collect Generated Images"
    hosts = ["comfyui"]
    families = ["image"]

    default_variant = "Main"

    def process(self, instance: pyblish.api.Instance):
        host: ComfyUIHost = registered_host()
        image_urls = host.stub.get_publish_node_images(instance.data)
        ext = ".png"
        instance.data["anatomyData"] = instance.context.data["anatomyData"]
        staging_dir = get_instance_staging_dir(instance)
        self.log.info("Outputting image to %s", staging_dir)

        files = []

        for image in image_urls:
            self.log.debug("Downloading image: %s", image)
            parse = urlsplit(image)
            self.log.debug(parse)
            query = parse_qs(parse.query)
            self.log.debug(query)
            filename = next(iter(query.get("filename")), None)
            if filename is None:
                continue
            self.log.debug("Got filename: %s", filename)
            destination = os.path.join(
                staging_dir, instance.data.get("productName"), filename
            )
            files.append(
                os.path.join(instance.data.get("productName"), filename)
            )
            Path(destination).parent.mkdir(parents=True, exist_ok=True)
            urlretrieve(image, destination)  # noqa: S310

        instance.context.data["currentFile"] = files[0]

        if len(files) == 1:
            files = files[0]

        # marking instance as reviewable
        instance.data["review"] = True
        instance.data["families"].append("review")

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

        # NOTE(@sas): Maybe generate a tiled image for batched images

        thumbnail_img = files
        if isinstance(thumbnail_img, list):
            thumbnail_img = thumbnail_img[0]

        # Thumbnail
        thumbnail = {
            "name": "thumbnail",
            "ext": ext[1:],
            "files": thumbnail_img,
            "stagingDir": staging_dir,
            "tags": ["thumbnail"],
        }

        instance.data["representations"].append(thumbnail)