Update app.py

app.py

@@ -1,120 +1,33 @@
 import gradio as gr
-import torch
-import json
-import time
 import subprocess
 import os
-from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Initialize DeepSeek
-model_name = "deepseek-ai/deepseek-coder-1.3b-instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-model.to("cpu")
-
-def generate_prompts_and_image(input_text):
-    start_time = time.time()
-
-    # Generate prompts with DeepSeek
-    prompt = f"""
-    Input: "{input_text}"
-    Task: Generate concise 'Positive' and 'Negative' AI image prompts for Stable Diffusion based on the input above. Output the prompts directly, no extra text or examples.
-    """
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-    outputs = model.generate(**inputs, max_new_tokens=50, temperature=0.7, top_p=0.9, do_sample=True)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-    if response.startswith(prompt):
-        response = response[len(prompt):].strip()
-
-    # Split response into positive and negative (assuming two lines)
-    lines = [line.strip() for line in response.split("\n") if line.strip()]
-    positive = lines[0] if lines else "No positive prompt generated"
-    negative = lines[1] if len(lines) > 1 else "No negative prompt generated"
-
-    # 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):
-        return {"error": "workflow.json not found", "time_taken": f"{time.time() - start_time:.2f} seconds"}
-
-    with open(workflow_path, "r") as f:
-        workflow = json.load(f)
-
-    # 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
+# Define the function to run hr.py
+def run_workflow():
     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_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 --input-directory or --output-directory directly; it reads from workflow
-        result = subprocess.run(
-            ["python", comfyui_main],
-            cwd=comfyui_dir,
-            capture_output=True,
-            text=True,
-            check=True
-        )
-        # ComfyUI saves images as ComfyUI_<timestamp>.png in output/
-        image_files = [f for f in os.listdir(output_dir) if f.startswith("ComfyUI") and f.endswith(".png")]
-        if not image_files:
-            image = f"ComfyUI ran but no image found: {result.stdout}\n{result.stderr}"
-        else:
-            image = os.path.join(output_dir, image_files[-1])  # Use the latest image
-    except subprocess.CalledProcessError as e:
-        image = f"ComfyUI failed: {e.stdout}\n{e.stderr}"
-
-    elapsed_time = time.time() - start_time
-    return {
-        "positive": positive,
-        "negative": negative,
-        "image": image,
-        "time_taken": f"{elapsed_time:.2f} seconds"
-    }
+        # Ensure the script exists
+        script_path = "ComfyUI/hr.py"
+        if not os.path.exists(script_path):
+            return "Error: hr.py not found!"
 
-def gradio_interface(input_text):
-    result = generate_prompts_and_image(input_text)
-    return result["image"], json.dumps({
-        "positive": result["positive"],
-        "negative": result["negative"],
-        "time_taken": result["time_taken"]
-    }, indent=2)
-
-demo = gr.Interface(
-    fn=gradio_interface,
-    inputs=gr.Textbox(label="Input Text", lines=10, placeholder="Paste your input text here..."),
-    outputs=[
-        gr.Image(label="Generated Image"),
-        gr.Textbox(label="Generated Prompts (JSON)")
-    ],
-    title="Prompt and Image Generator"
+        # Run the script
+        result = subprocess.run(["python", script_path], capture_output=True, text=True)
+        
+        # Return the output
+        return result.stdout if result.returncode == 0 else f"Error:\n{result.stderr}"
+    
+    except Exception as e:
+        return f"Exception: {str(e)}"
+
+# Create a simple UI with a button
+ui = gr.Interface(
+    fn=run_workflow,
+    inputs=[],
+    outputs="text",
+    title="Run ComfyUI Workflow",
+    description="Click the button to execute hr.py and generate an image."
 )
 
+# Launch the Gradio app
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    ui.launch()