Update app.py

app.py

@@ -9,38 +9,36 @@ def extract_metadata(image):
 
     try:
         metadata = {}
-        if 'metadata' in image.info:
-            metadata = json.loads(image.info['metadata'])
-        elif 'prompt' in image.info:
-            metadata = json.loads(image.info['prompt'])
-        elif 'Comment' in image.info:
-            metadata = json.loads(image.info['Comment'])
-            metadata['model'] = 'NovelAI'
-        elif 'parameters' in image.info:
-            if image.info['parameters'].startswith('{'):
-                parameters_data = json.loads(image.info['parameters'])
-                if 'sui_image_params' in parameters_data:
-                    sui_image_params = parameters_data['sui_image_params']
-                    metadata.update(sui_image_params)
-                else:
-                    metadata = parameters_data
-            else:
-                lines = image.info['parameters'].split('\n')
-                prompt = lines[0].strip()
-                negative_prompt = lines[1].strip().replace('Negative prompt:', '').strip()
-                metadata['prompt'] = prompt
-                metadata['negative_prompt'] = negative_prompt
+        # Handling multiple possible metadata keys
+        potential_keys = ['metadata', 'prompt', 'Comment', 'parameters', 'exif']
+        for key in potential_keys:
+            if key in image.info:
+                raw_data = image.info[key]
                 
-                for line in lines[2:]:
-                    line = line.strip()
-                    if line.startswith('Steps:'):
-                        steps_info = line.split(':', 1)[1].strip().split(',')
-                        for info in steps_info:
-                            info = info.strip()
-                            if ':' in info:
-                                key, value = info.split(':', 1)
+                # If raw data starts with '{', assume JSON format
+                if raw_data.startswith('{'):
+                    metadata = json.loads(raw_data)
+                else:
+                    if key == 'parameters':
+                        # Attempt to process Stable Diffusion or NovelAI style data
+                        lines = raw_data.split('\n')
+                        prompt = lines[0].strip()
+                        negative_prompt = lines[1].strip().replace('Negative prompt:', '').strip()
+                        metadata['prompt'] = prompt
+                        metadata['negative_prompt'] = negative_prompt
+
+                        for line in lines[2:]:
+                            line = line.strip()
+                            if ':' in line:
+                                key, value = line.split(':', 1)
                                 metadata[key.strip()] = value.strip()
-        else:
+                    elif key == 'Comment':
+                        # Specific handling for NovelAI
+                        metadata = json.loads(raw_data)
+                        metadata['model'] = 'NovelAI'
+                break  # Exit loop once a supported key is found
+
+        if not metadata:
             return "No supported metadata found in the image.", {}
 
         return "Metadata extracted successfully.", metadata
@@ -68,6 +66,7 @@ with gr.Blocks() as demo:
             status_output = gr.Textbox(label="Status")
             output_metadata = gr.JSON(label="Metadata")
 
+    # Event listener for when the image is changed
     input_image.change(
         fn=process_image,
         inputs=input_image,
@@ -75,6 +74,4 @@ with gr.Blocks() as demo:
         api_name="interrogate"
     )
 
-    
-
-demo.launch()
+demo.launch()