v3p1

utils.py

@@ -6,8 +6,7 @@ import json
 import torch
 from PIL import Image, PngImagePlugin
 from datetime import datetime
-from dataclasses import dataclass
-from typing import Callable, Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple
 from diffusers import (
     DDIMScheduler,
     DPMSolverMultistepScheduler,
@@ -15,8 +14,6 @@ from diffusers import (
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
 )
-import base64
-from io import BytesIO
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -46,7 +43,7 @@ def aspect_ratio_handler(aspect_ratio: str, custom_width: int, custom_height: in
         width, height = parse_aspect_ratio(aspect_ratio)
         return width, height
 
-def get_scheduler(scheduler_config: Dict, name: str) -> Optional[Callable]:
+def get_scheduler(scheduler_config: Dict, name: str) -> Optional[DDIMScheduler]:
     scheduler_factory_map = {
         "DPM++ 2M Karras": lambda: DPMSolverMultistepScheduler.from_config(scheduler_config, use_karras_sigmas=True),
         "DPM++ SDE Karras": lambda: DPMSolverSinglestepScheduler.from_config(scheduler_config, use_karras_sigmas=True),
@@ -95,6 +92,15 @@ def is_google_colab():
     except:
         return False
 
+def get_image_metadata(image_path):
+    try:
+        with Image.open(image_path) as img:
+            metadata = img.info.get('metadata', '{}')
+            return json.loads(metadata)
+    except Exception as e:
+        print(f"Error reading metadata from {image_path}: {str(e)}")
+        return {}
+
 def validate_json_parameters(json_str):
     try:
         params = json.loads(json_str)
@@ -106,9 +112,4 @@ def validate_json_parameters(json_str):
     except json.JSONDecodeError:
         raise ValueError("Invalid JSON format")
     except Exception as e:
-        raise ValueError(f"Error parsing JSON: {str(e)}")
-
-def image_to_base64(image):
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+        raise ValueError(f"Error parsing JSON: {str(e)}")