Upload multi_comfy.py

modules/multi_comfy.py

@@ -0,0 +1,384 @@
+import gradio as gr
+import json
+from PIL import Image
+import os
+from collections import defaultdict
+
+css = """
+#custom-gallery{--row-height:180px;display:grid;grid-auto-rows:min-content;gap:10px}#custom-gallery .thumbnail-item{height:var(--row-height);width:100%;position:relative;overflow:hidden;border-radius:8px;box-shadow:0 2px 5px rgb(0 0 0 / .1);transition:transform 0.2s ease,box-shadow 0.2s ease}#custom-gallery .thumbnail-item:hover{transform:translateY(-3px);box-shadow:0 4px 12px rgb(0 0 0 / .15)}#custom-gallery .thumbnail-item img{width:auto;height:100%;max-width:100%;max-height:var(--row-height);object-fit:contain;margin:0 auto;display:block}#custom-gallery .thumbnail-item img.portrait{max-width:100%}#custom-gallery .thumbnail-item img.landscape{max-height:100%}.gallery-container{max-height:500px;overflow-y:auto;padding-right:0;--size-80:500px}.thumbnails{display:flex;position:absolute;bottom:0;width:120px;overflow-x:scroll;padding-top:320px;padding-bottom:280px;padding-left:4px;flex-wrap:wrap}
+"""
+
+EMPTY_RESULT = ("Not Available",) * 15
+
+# ---------- EXTRACTION FUNCTIONS ----------
+def read_metadata(file_path):
+    try:
+        with Image.open(file_path) as img:
+            return img.info
+    except Exception as e:
+        return {"error": f"Error reading file: {str(e)}"}
+
+def extract_workflow_data(file_path):
+    metadata = read_metadata(file_path)
+    if "error" in metadata:
+        return {"error": metadata["error"]}
+
+    if 'prompt' in metadata:
+        try:
+            return json.loads(metadata['prompt'])
+        except json.JSONDecodeError:
+            pass
+
+    for key, value in metadata.items():
+        if isinstance(value, str) and value.strip().startswith('{'):
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError:
+                continue
+    return {"error": "No workflow data found"}
+
+def extract_ksampler_params(workflow_data):
+    seed = steps = cfg = sampler = scheduler = denoise = "Not found"
+    if not isinstance(workflow_data, dict):
+        return seed, steps, cfg, sampler, scheduler, denoise
+    for node in workflow_data.values():
+        if isinstance(node, dict) and node.get("class_type", "") in ["KSampler", "KSampler (Efficient)"]:
+            inputs = node.get("inputs", {})
+            seed = inputs.get("seed", "Not found")
+            steps = inputs.get("steps", "Not found")
+            cfg = inputs.get("cfg", "Not found")
+            sampler = inputs.get("sampler_name", "Not found")
+            scheduler = inputs.get("scheduler", "Not found")
+            denoise = inputs.get("denoise", "Not found")
+            break
+    return str(seed), str(steps), str(cfg), str(sampler), str(scheduler), str(denoise)
+
+def extract_prompts(workflow_data):
+    positive = negative = "Not found"
+    if not isinstance(workflow_data, dict):
+        return positive, negative
+    for node in workflow_data.values():
+        if isinstance(node, dict):
+            class_type = node.get("class_type", "")
+            inputs = node.get("inputs", {})
+            title = node.get("_meta", {}).get("title", "") if node.get("_meta") else ""
+
+            if "Text to Conditioning" in class_type:
+                if "POSITIVE" in title:
+                    positive = inputs.get("text", "Not found")
+                elif "NEGATIVE" in title:
+                    negative = inputs.get("text", "Not found")
+            if "ShowText|pysssss" in class_type:
+                if "text_1" in inputs:
+                    positive = inputs["text_1"]
+                if "text_2" in inputs:
+                    negative = inputs["text_2"]
+            if "DPRandomGenerator" in class_type:
+                if "POSITIVE" in title:
+                    positive = inputs.get("text", "Not found")
+                elif "NEGATIVE" in title:
+                    negative = inputs.get("text", "Not found")
+    return str(positive), str(negative)
+
+def extract_loras(workflow_data):
+    loras = []
+    if not isinstance(workflow_data, dict):
+        return "None found"
+    for node in workflow_data.values():
+        if isinstance(node, dict):
+            inputs = node.get("inputs", {})
+            if "LoraLoader" in node.get("class_type", ""):
+                name = inputs.get("lora_name", "Unknown")
+                strength = inputs.get("strength_model", "Unknown")
+                loras.append(f"{name} (Strength: {strength})")
+            for val in inputs.values():
+                if isinstance(val, str) and "lora:" in val.lower():
+                    loras.append(val)
+    return "\n".join(loras) if loras else "None found"
+
+def extract_model_info(workflow_data):
+    models = []
+    if not isinstance(workflow_data, dict):
+        return "Not found"
+    for node in workflow_data.values():
+        if isinstance(node, dict):
+            inputs = node.get("inputs", {})
+            class_type = node.get("class_type", "")
+            if "CheckpointLoader" in class_type:
+                models.append(inputs.get("ckpt_name", "Unknown"))
+            if "Model Mecha Recipe" in class_type:
+                models.append(inputs.get("model_path", "Unknown"))
+    return "\n".join(models) if models else "Not found"
+
+def extract_image_info_from_file(image_path):
+    """Extract actual image dimensions from the image file itself"""
+    try:
+        with Image.open(image_path) as img:
+            width, height = img.size
+            return str(width), str(height)
+    except Exception as e:
+        return "Not found", "Not found"
+
+def extract_batch_size(workflow_data):
+    """Extract batch size from workflow data"""
+    batch_size = "Not found"
+    if not isinstance(workflow_data, dict):
+        return batch_size
+    for node in workflow_data.values():
+        if isinstance(node, dict) and node.get("class_type", "") == "EmptyLatentImage":
+            inputs = node.get("inputs", {})
+            batch_size = inputs.get("batch_size", "Not found")
+            break
+    return str(batch_size)
+
+def extract_nodes_info(workflow_data):
+    if not isinstance(workflow_data, dict):
+        return "Not found"
+    total_nodes = len(workflow_data)
+    node_types = defaultdict(int)
+    for node in workflow_data.values():
+        if isinstance(node, dict):
+            node_types[node.get("class_type", "Unknown")] += 1
+    summary = f"Total Nodes: {total_nodes}\n"
+    for t, c in sorted(node_types.items()):
+        summary += f"{t}: {c}\n"
+    return summary.strip()
+
+def extract_workflow_as_json(workflow_data):
+    if isinstance(workflow_data, dict):
+        return json.dumps(workflow_data, ensure_ascii=False, indent=2)
+    return "{}"
+# ---------- EXTRACTION FUNCTIONS ----------
+#
+# ---------- IMAGE PROCESSING ----------
+def process_single_image(image_path):
+    """Extract all workflow info from a single image path."""
+    if not image_path:
+        return EMPTY_RESULT
+
+    workflow_data = extract_workflow_data(image_path)
+
+    if isinstance(workflow_data, dict) and "error" not in workflow_data:
+        seed, steps, cfg, sampler, scheduler, denoise = extract_ksampler_params(workflow_data)
+        positive, negative = extract_prompts(workflow_data)
+        loras = extract_loras(workflow_data)
+        models = extract_model_info(workflow_data)
+
+        # Get actual image dimensions instead of workflow dimensions
+        width, height = extract_image_info_from_file(image_path)
+        batch = extract_batch_size(workflow_data)
+
+        nodes = extract_nodes_info(workflow_data)
+        full_json = extract_workflow_as_json(workflow_data)
+    else:
+        error = str(workflow_data.get("error", "Unknown error"))
+        seed = steps = cfg = sampler = scheduler = denoise = positive = negative = loras = models = width = height = batch = nodes = full_json = error
+
+    return seed, steps, cfg, sampler, scheduler, denoise, \
+           positive, negative, loras, models, width, height, batch, nodes, full_json
+
+def append_gallery(gallery: list, image: str):
+    """Add a single image to the gallery"""
+    if gallery is None:
+        gallery = []
+    if not image:
+        return gallery, None
+    gallery.append(image)
+    return gallery, None
+
+def extend_gallery(gallery, images):
+    """Extend gallery preserving uniqueness"""
+
+    if gallery is None:
+        gallery = []
+
+    if not images:
+        return gallery
+
+    # Normalize input - Gradio might pass various formats
+    incoming_paths = []
+    if isinstance(images, str):  # Single image path
+        incoming_paths.append(images)
+    elif isinstance(images, list):
+        for img in images:
+            # Handle cases where elements could be tuples from Gallery
+            if isinstance(img, (tuple, list)):
+                incoming_paths.append(str(img[0]))
+            else:
+                incoming_paths.append(str(img))
+
+    unique_incoming = list(set(incoming_paths))  # Avoid duplicates
+
+    seen_paths = {item[0] if isinstance(item, (list, tuple)) else item for item in gallery}
+    new_entries = [path for path in unique_incoming if path not in seen_paths]
+
+    # Create entries matching expected gallery style
+    formatted_new = [(path, '') for path in new_entries]
+
+    updated_gallery = gallery + formatted_new
+
+    return updated_gallery
+
+def process_gallery(gallery, results_state):
+    """Process all images and populate metadata in session."""
+    if not gallery or len(gallery) == 0:
+        # Clear results if nothing left
+        results_state.clear()
+        return EMPTY_RESULT + (results_state,)
+
+    updated_state = {}
+    first_image_result = EMPTY_RESULT
+    try:
+        for item in gallery:
+            path = item if isinstance(item, str) else item[0]
+
+            if path not in results_state:
+                res = process_single_image(path)
+                results_state[path] = res
+                updated_state[path] = res
+
+                if first_image_result == EMPTY_RESULT:
+                    first_image_result = res
+            else:
+                # Already cached
+                res = results_state[path]
+                updated_state[path] = res
+
+                if first_image_result == EMPTY_RESULT:
+                    first_image_result = res
+
+        results_state.update(updated_state)
+        return first_image_result + (results_state,)
+    except Exception as e:
+        print("[ERROR]", str(e))
+        return EMPTY_RESULT + (results_state,)
+
+def get_selection_from_gallery(gallery, results_state, evt: gr.SelectData):
+    """Fetch result for selected image in gallery."""
+    if evt is None or evt.value is None:
+        # No selection: use first image
+        if gallery and len(gallery) > 0:
+            img_path = str(gallery[0][0] if isinstance(gallery[0], (list, tuple)) else gallery[0])
+            if img_path in results_state:
+                return list(results_state[img_path])
+    else:
+        # Handle selection event
+        try:
+            selected_value = evt.value
+            img_path = None
+
+            if isinstance(selected_value, dict) and 'image' in selected_value:
+                img_path = selected_value['image']['path']
+            elif isinstance(selected_value, (list, tuple)):
+                img_path = selected_value[0]
+            else:
+                img_path = str(selected_value)
+
+            if img_path in results_state:
+                return list(results_state[img_path])
+        except Exception as e:
+            print(f"Selection error: {e}")
+
+    # Return empty if no image found
+    return list(EMPTY_RESULT)
+# ---------- IMAGE PROCESSING ----------
+#
+def create_multi_comfy():
+    with gr.Blocks(css=css, fill_width=True) as demo:
+        gr.Markdown("# 🛠️ ComfyUI Workflow Information Extractor")
+        gr.Markdown("Upload Multiple ComfyUI-generated images. Extract prompts, parameters, models, and full workflows.")
+        with gr.Row():
+            with gr.Column(scale=2):
+                upload_button = gr.UploadButton(
+                    "📁 Upload Multiple Images",
+                    file_types=["image"],
+                    file_count="multiple",
+                    size='lg'
+                )
+                gallery = gr.Gallery(
+                    columns=3,
+                    show_share_button=False,
+                    interactive=True,
+                    height='auto',
+                    label='Grid of images',
+                    preview=False,
+                    elem_id='custom-gallery'
+                )
+            with gr.Column(scale=3):
+                with gr.Tabs():
+                    with gr.Tab("Sampling Parameters"):
+                        with gr.Row():
+                            with gr.Column():
+                                seed_out = gr.Textbox(label="Seed", interactive=False, show_copy_button=True)
+                                steps_out = gr.Textbox(label="Steps", interactive=False, show_copy_button=True)
+                                cfg_out = gr.Textbox(label="CFG Scale", interactive=False)
+                            with gr.Column():
+                                sampler_out = gr.Textbox(label="Sampler", interactive=False)
+                                scheduler_out = gr.Textbox(label="Scheduler", interactive=False)
+                                denoise_out = gr.Textbox(label="Denoise", interactive=False)
+
+                    with gr.Tab("Prompts"):
+                        pos_prompt = gr.Textbox(label="Positive Prompt", lines=4, interactive=False, show_copy_button=True)
+                        neg_prompt = gr.Textbox(label="Negative Prompt", lines=4, interactive=False, show_copy_button=True)
+
+                    with gr.Tab("Models & LoRAs"):
+                        with gr.Row():
+                            lora_out = gr.Textbox(label="LoRAs", lines=5, interactive=False, show_copy_button=True)
+                            model_out = gr.Textbox(label="Base Models", lines=5, interactive=False, show_copy_button=True)
+
+                    with gr.Tab("Image Info"):
+                        with gr.Row():
+                            with gr.Column():
+                                width_out = gr.Textbox(label="Width", interactive=False)
+                                height_out = gr.Textbox(label="Height", interactive=False)
+                                batch_out = gr.Textbox(label="Batch Size", interactive=False)
+                            with gr.Column():
+                                nodes_out = gr.Textbox(label="Node Counts", lines=15, interactive=True, show_copy_button=True)
+
+                    with gr.Tab("Full Workflow"):
+                        json_out = gr.Textbox(label="Workflow JSON", lines=20, interactive=True, show_copy_button=True)
+
+        # State to store results per image
+        results_state = gr.State({})
+
+        # Event Connections
+        upload_event = upload_button.upload(
+            fn=extend_gallery,
+            inputs=[gallery, upload_button],
+            outputs=gallery,
+            queue=False
+        )
+
+        upload_event.then(
+            fn=process_gallery,
+            inputs=[gallery, results_state],
+            outputs=[
+                seed_out, steps_out, cfg_out, sampler_out, scheduler_out, denoise_out,
+                pos_prompt, neg_prompt, lora_out, model_out, width_out, height_out,
+                batch_out, nodes_out, json_out, results_state
+            ]
+        )
+        gallery.change(
+            fn=process_gallery,
+            inputs=[gallery, results_state],
+            outputs=[
+                seed_out, steps_out, cfg_out, sampler_out, scheduler_out, denoise_out,
+                pos_prompt, neg_prompt, lora_out, model_out, width_out, height_out,
+                batch_out, nodes_out, json_out, results_state
+            ],
+            queue=True
+        )
+
+        gallery.select(
+            get_selection_from_gallery,
+            inputs=[gallery, results_state],
+            outputs=[
+                seed_out, steps_out, cfg_out, sampler_out, scheduler_out, denoise_out,
+                pos_prompt, neg_prompt, lora_out, model_out, width_out, height_out,
+                batch_out, nodes_out, json_out
+            ]
+        )
+
+        gr.Markdown("---\n💡 **Note:** It's under development.")
+    
+    return demo