File size: 4,026 Bytes
34b6c63 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
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"
# Load and modify 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)
# 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)
# 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
result = subprocess.run(
["python", comfyui_main, "--input-directory", comfyui_dir, "--output-directory", output_dir],
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"
}
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"
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860) |