Update app.py

app.py

@@ -31,7 +31,15 @@ def generate_prompts_and_image(input_text):
     positive = lines[0] if lines else "No positive prompt generated"
     negative = lines[1] if len(lines) > 1 else "No negative prompt generated"
 
-    # Load and modify ComfyUI workflow
+    # Debug: Check if ComfyUI and model exist
+    comfyui_dir = os.path.join(os.path.dirname(__file__), "ComfyUI")
+    checkpoint_path = os.path.join(comfyui_dir, "models", "checkpoints", "v1-5-pruned-emaonly-fp16.safetensors")
+    if not os.path.exists(comfyui_dir):
+        return {"error": f"ComfyUI directory not found at {comfyui_dir}", "time_taken": f"{time.time() - start_time:.2f} seconds"}
+    if not os.path.exists(checkpoint_path):
+        return {"error": f"Checkpoint not found at {checkpoint_path}", "time_taken": f"{time.time() - start_time:.2f} seconds"}
+
+    # Load and debug ComfyUI workflow
     workflow_path = "workflow.json"
     temp_workflow_path = "temp_workflow.json"
     if not os.path.exists(workflow_path):
@@ -40,24 +48,34 @@ def generate_prompts_and_image(input_text):
     with open(workflow_path, "r") as f:
         workflow = json.load(f)
 
-    # Inject prompts into nodes 6 (positive) and 7 (negative) from your workflow.json
-    workflow["6"]["widgets_values"][0] = positive  # Positive prompt (CLIPTextEncode)
-    workflow["7"]["widgets_values"][0] = negative  # Negative prompt (CLIPTextEncode)
+    # Debug: Print workflow nodes to verify IDs
+    print("Workflow nodes:", json.dumps(workflow["nodes"], indent=2))
+
+    # Use hardcoded node IDs from workflow.json (6 for positive, 7 for negative)
+    positive_id, negative_id = 6, 7  # Hardcoded based on your workflow.json
+    try:
+        if str(positive_id) not in workflow or str(negative_id) not in workflow:
+            return {"error": f"Node IDs {positive_id} or {negative_id} not found in workflow.json", "time_taken": f"{time.time() - start_time:.2f} seconds"}
+        
+        # Inject prompts into nodes 6 (positive) and 7 (negative)
+        workflow[str(positive_id)]["widgets_values"][0] = positive  # Positive prompt (CLIPTextEncode)
+        workflow[str(negative_id)]["widgets_values"][0] = negative  # Negative prompt (CLIPTextEncode)
+    except KeyError as e:
+        return {"error": f"Error injecting prompts: {str(e)}", "time_taken": f"{time.time() - start_time:.2f} seconds"}
 
     # Save temporary workflow
     with open(temp_workflow_path, "w") as f:
         json.dump(workflow, f)
 
     # Run ComfyUI via subprocess
-    comfyui_dir = os.path.join(os.path.dirname(__file__), "ComfyUI")
     comfyui_main = os.path.join(comfyui_dir, "main.py")
     output_dir = os.path.join(comfyui_dir, "output")  # ComfyUI default output dir
     os.makedirs(output_dir, exist_ok=True)
 
     try:
-        # Note: ComfyUI's main.py doesn't use --workflow directly; it reads from the workflow file
+        # Note: ComfyUI's main.py doesn't use --input-directory or --output-directory directly; it reads from workflow
         result = subprocess.run(
-            ["python", comfyui_main, "--input-directory", comfyui_dir, "--output-directory", output_dir],
+            ["python", comfyui_main],
             cwd=comfyui_dir,
             capture_output=True,
             text=True,