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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143 | class CollectVideo(pyblish.api.InstancePlugin):
"""Collect video for publish."""
order = pyblish.api.CollectorOrder + 0.16
label = "Collect Generated Video + Thumbnail"
hosts = ["comfyui"]
families = ["render"]
default_variant = "Main"
def process(self, instance: pyblish.api.Instance):
host: ComfyUIHost = registered_host()
image_urls = host.stub.get_publish_node_images(
instance.data, publish_type=PublishType.VIDEO
)
video_info = VideoInfo()
video_exts = {".mp4", ".webm"}
video_info.thumbnail_extension = ".png"
# f"{filename}_{format}_thumb.png"
instance.data["anatomyData"] = instance.context.data["anatomyData"]
staging_dir = get_instance_staging_dir(instance)
self.log.debug("Outputting video to %s", staging_dir)
video_link = next(iter(image_urls))
if video_link is None:
self.log.warning("Nothing could be collected. (No url returned.)")
return
# Download video
self.log.debug(video_link)
parse = urlsplit(video_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 video_exts:
self.log.warning(
"Nothing could be collected. "
"(filename has invalid extension for video.)"
)
return
video_info.video_extension = extension
self.log.debug(filename)
self.log.debug(staging_dir)
destination = os.path.join(
staging_dir, instance.data.get("productName"), filename
)
video_info.video_file = os.path.join(
instance.data.get("productName"), filename
)
Path(destination).parent.mkdir(parents=True, exist_ok=True)
urlretrieve(video_link, destination) # noqa: S310
# retrieve thumbnail
thumb_filename = f"{Path(filename).stem}_{extension[1:]}_thumb.png"
query["filename"] = [thumb_filename]
thumb_url = parse._replace(
query=naive_reconstruct_querydict(query)
).geturl()
self.log.debug("Retrieving generated thumbnail")
self.log.debug(thumb_url)
thumb_destination = os.path.join(
staging_dir, instance.data.get("productName"), thumb_filename
)
urlretrieve(thumb_url, thumb_destination) # noqa: S310
video_info.thumbnail_file = os.path.join(
instance.data.get("productName"), thumb_filename
)
instance.context.data["currentFile"] = video_info.video_file
# marking instance as reviewable
instance.data["review"] = True
instance.data["families"].append("review")
# creating representation
instance.data["representations"].append(
{
"name": video_info.video_extension[1:],
"ext": video_info.video_extension[1:],
"files": video_info.video_file,
"stagingDir": staging_dir,
"tags": ["review"],
}
)
# Thumbnail
thumbnail = {
"name": "thumbnail",
"ext": video_info.thumbnail_extension[1:],
"files": video_info.thumbnail_file,
"stagingDir": staging_dir,
"tags": ["thumbnail"],
}
instance.data["representations"].append(thumbnail)
|