get mimetype from buffer with no descriptor - issue Stability-AI#125

gacogo · Nov 19, 2022 · 72882fa · 72882fa
1 parent 67b5cf2
commit 72882fa
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/stability_sdk/client.py b/src/stability_sdk/client.py
@@ -11,6 +11,7 @@
 import logging
 import time
 import mimetypes
+import magic
 
 import grpc
 from argparse import ArgumentParser, Namespace
@@ -86,7 +87,8 @@ def process_artifacts_from_answers(
  for artifact in resp.artifacts:
  artifact_start = time.time()
  if artifact.type == generation.ARTIFACT_IMAGE:
- ext = mimetypes.guess_extension(artifact.mime)
+ mime_type = magic.from_buffer(artifact.binary, mime=True)
+ ext = mime_type.split("/")[1]
  contents = artifact.binary
  elif artifact.type == generation.ARTIFACT_CLASSIFICATIONS:
  ext = ".pb.json"