152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
import os
|
|
import json
|
|
import time
|
|
import requests
|
|
import copy
|
|
import subprocess
|
|
|
|
COMFYUI_URL = "http://127.0.0.1:8000"
|
|
OUTPUT_DIR = "/Users/davidkotnik/repos/novafarma/assets/images"
|
|
|
|
workflow = {
|
|
"3": {
|
|
"inputs": {
|
|
"seed": int(time.time()),
|
|
"steps": 25,
|
|
"cfg": 7.5,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 1,
|
|
"model": ["4", 0],
|
|
"positive": ["6", 0],
|
|
"negative": ["7", 0],
|
|
"latent_image": ["5", 0]
|
|
},
|
|
"class_type": "KSampler"
|
|
},
|
|
"4": {
|
|
"inputs": {
|
|
"ckpt_name": "v1-5-pruned-emaonly.safetensors"
|
|
},
|
|
"class_type": "CheckpointLoaderSimple"
|
|
},
|
|
"5": {
|
|
"inputs": {
|
|
"width": 512,
|
|
"height": 512,
|
|
"batch_size": 1
|
|
},
|
|
"class_type": "EmptyLatentImage"
|
|
},
|
|
"6": {
|
|
"inputs": {
|
|
"text": "",
|
|
"clip": ["4", 1],
|
|
"class_type": "CLIPTextEncode"
|
|
},
|
|
"7": {
|
|
"inputs": {
|
|
"text": "blurry, low quality, pixelated, photo, realistic",
|
|
"clip": ["4", 1]
|
|
},
|
|
"class_type": "CLIPTextEncode"
|
|
},
|
|
"8": {
|
|
"inputs": {
|
|
"samples": ["3", 0],
|
|
"vae": ["4", 2]
|
|
},
|
|
"class_type": "VAEDecode"
|
|
},
|
|
"9": {
|
|
"inputs": {
|
|
"filename_prefix": "",
|
|
"images": ["8", 0],
|
|
},
|
|
"class_type": "SaveImage"
|
|
}
|
|
}
|
|
|
|
def queue_prompt(prompt):
|
|
p = {"prompt": prompt}
|
|
data = json.dumps(p).encode('utf-8')
|
|
try:
|
|
res = requests.post(COMFYUI_URL + "/prompt", data=data)
|
|
if res.status_code == 200:
|
|
print("🚀 SUCCESS: Poslano v ComfyUI!")
|
|
return res.json()
|
|
else:
|
|
print(f"❌ Napaka: {res.text}")
|
|
return None
|
|
except Exception as e:
|
|
print(f"❌ UFI (ComfyUI) ni prižgan! {e}")
|
|
return None
|
|
|
|
def wait_for_completion(prompt_id: str, timeout: int = 180) -> bool:
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
try:
|
|
r = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
|
|
history = r.json()
|
|
if prompt_id in history:
|
|
status = history[prompt_id].get('status', {})
|
|
if status.get('completed'):
|
|
return True
|
|
except Exception:
|
|
pass
|
|
time.sleep(2)
|
|
return False
|
|
|
|
def download_image(prompt_id: str, prefix: str):
|
|
try:
|
|
r = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
|
|
history = r.json()
|
|
if prompt_id not in history:
|
|
return False
|
|
outputs = history[prompt_id].get('outputs', {})
|
|
for node_id, out in outputs.items():
|
|
if 'images' in out:
|
|
for img in out['images']:
|
|
filename = img['filename']
|
|
subfolder = img.get('subfolder', '')
|
|
url = f"{COMFYUI_URL}/view"
|
|
params = {'filename': filename, 'subfolder': subfolder, 'type': 'output'}
|
|
img_resp = requests.get(url, params=params)
|
|
if img_resp.status_code == 200:
|
|
out_path = os.path.join(OUTPUT_DIR, f"{prefix}_{filename}")
|
|
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
|
with open(out_path, 'wb') as f:
|
|
f.write(img_resp.content)
|
|
print(f"✅ Image saved: {out_path}")
|
|
# subprocess.run(["open", out_path])
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
print(f"❌ Download error: {e}")
|
|
return False
|
|
|
|
# Updated generate_character to use new helpers
|
|
|
|
def generate_character(name, prompt_text, prefix):
|
|
wf = copy.deepcopy(workflow)
|
|
wf["6"]["inputs"]["text"] = prompt_text
|
|
wf["9"]["inputs"]["filename_prefix"] = prefix
|
|
result = queue_prompt(wf)
|
|
if not result or 'prompt_id' not in result:
|
|
print('❌ Failed to get prompt_id')
|
|
return
|
|
pid = result['prompt_id']
|
|
if wait_for_completion(pid):
|
|
download_image(pid, prefix)
|
|
else:
|
|
print('❌ Generation timed out')
|
|
|
|
# Generate 5 Kai images
|
|
kai_prompt = "Top-down 2.5D, (bold black outlines:1.4), dark hand-drawn stylized indie game art, high contrast noir style, character Kai with pink and green dreadlocks, friendly smile"
|
|
for i in range(5):
|
|
generate_character("Kai", kai_prompt, f"noir_kai_{i}")
|
|
|
|
# Generate 5 Gronk images
|
|
gronk_prompt = "Top-down 2.5D, (bold black outlines:1.4), dark hand-drawn stylized indie game art, high contrast noir style, character Gronk with pink dreadlocks, stretched ears, facial piercings, holding a vape"
|
|
for i in range(5):
|
|
generate_character("Gronk", gronk_prompt, f"noir_gronk_{i}") |