Michael Yang
commited on
Commit
·
811f626
1
Parent(s):
28909fe
convert from numpy
Browse files- baseline.py +14 -4
baseline.py
CHANGED
|
@@ -40,12 +40,22 @@ def run(prompt, scheduler_key='dpm_scheduler', bg_seed=1, num_inference_steps=20
|
|
| 40 |
model_dict, latents, input_embeddings, num_inference_steps,
|
| 41 |
guidance_scale=guidance_scale, scheduler_key=scheduler_key
|
| 42 |
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 43 |
|
| 44 |
gc.collect()
|
| 45 |
torch.cuda.empty_cache()
|
| 46 |
|
| 47 |
-
with BytesIO() as buffer:
|
| 48 |
-
np.save(buffer, images[0])
|
| 49 |
-
img_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
|
| 50 |
|
| 51 |
-
return images[0],
|
|
|
|
| 40 |
model_dict, latents, input_embeddings, num_inference_steps,
|
| 41 |
guidance_scale=guidance_scale, scheduler_key=scheduler_key
|
| 42 |
)
|
| 43 |
+
|
| 44 |
+
# Convert to PIL Image
|
| 45 |
+
import PIL.Image
|
| 46 |
+
image = PIL.Image.fromarray(images[0])
|
| 47 |
+
|
| 48 |
+
# Save as PNG in memory
|
| 49 |
+
buffer = BytesIO()
|
| 50 |
+
image.save(buffer, format='PNG')
|
| 51 |
+
|
| 52 |
+
# Encode PNG to base64
|
| 53 |
+
png_bytes = buffer.getvalue()
|
| 54 |
+
base64_string = base64.b64encode(png_bytes).decode('utf-8')
|
| 55 |
+
|
| 56 |
|
| 57 |
gc.collect()
|
| 58 |
torch.cuda.empty_cache()
|
| 59 |
|
|
|
|
|
|
|
|
|
|
| 60 |
|
| 61 |
+
return images[0], base64_string
|