opencode/generate-image.py at dev · Grigory-Rylov/opencode · GitHub

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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
#!/usr/bin/env python3

import requests
import json
import sys
import time
import os
from pathlib import Path

COMFYUI_HOST = os.environ.get('COMFYUI_HOST', '192.168.1.212')
COMFYUI_PORT = os.environ.get('COMFYUI_PORT', '8188')
WORKFLOW_FILE = os.environ.get('WORKFLOW_FILE', os.path.expanduser('~/flux2_txt2img.json'))

def check_comfyui():
    try:
        requests.get(f'http://{COMFYUI_HOST}:{COMFYUI_PORT}/', timeout=5)
        return True
    except:
        return False

def queue_prompt(prompt, seed=None):
    with open(WORKFLOW_FILE, 'r') as f:
        workflow = json.load(f)

    # Replace empty text with prompt
    for node_id, node in workflow.items():
        if 'inputs' in node and 'text' in node['inputs']:
            if node['inputs']['text'] == '':
                node['inputs']['text'] = prompt
                break

    if seed is not None:
        workflow['25']['inputs']['noise_seed'] = seed

    response = requests.post(
        f'http://{COMFYUI_HOST}:{COMFYUI_PORT}/prompt',
        json={'prompt': workflow},
        headers={'Content-Type': 'application/json'}
    )

    return response.json()['prompt_id']

def wait_for_completion(prompt_id):
    start_time = time.time()
    timeout = 5 * 60  # 5 minutes

    while time.time() - start_time < timeout:
        try:
            queue_response = requests.get(f'http://{COMFYUI_HOST}:{COMFYUI_PORT}/queue')
            queue_data = queue_response.json()

            if any(prompt_id in r for r in queue_data.get('queue_running', [])):
                time.sleep(5)
                continue

            history_response = requests.get(f'http://{COMFYUI_HOST}:{COMFYUI_PORT}/history/{prompt_id}')
            history = history_response.json()

            if history and len(history) > 0:
                return history[prompt_id]

            time.sleep(5)
        except:
            time.sleep(5)

    raise Exception('Generation timed out after 5 minutes')

def find_output_image():
    output_dir = Path('/home/grishberg/ComfyUI/output')

    if not output_dir.exists():
        return None

    flux_files = sorted(
        output_dir.glob('flux2_gen_*.png'),
        key=lambda p: p.stat().st_mtime,
        reverse=True
    )

    return flux_files[0] if flux_files else None

def generate_image(prompt, output=None, seed=None):
    if not check_comfyui():
        raise Exception(f'ComfyUI is not running at {COMFYUI_HOST}:{COMFYUI_PORT}')

    print(f'Generating image... (seed: {seed or "random"})')

    prompt_id = queue_prompt(prompt, seed)
    print(f'Queue: {prompt_id}')

    result = wait_for_completion(prompt_id)

    output_path = find_output_image()

    if not output_path:
        raise Exception('Output image not found')

    if output:
        import shutil
        shutil.copy2(str(output_path), output)
        print(f'Saved to: {output}')
    else:
        print(f'Output: {output_path}')

    return str(output_path)

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print('Usage: generate-image.py <prompt> <output.png> [seed]')
        sys.exit(1)

    prompt = sys.argv[1]
    output = sys.argv[2]
    seed = int(sys.argv[3]) if len(sys.argv) > 3 else None

    try:
        generate_image(prompt, output, seed)
    except Exception as e:
        print(f'Error: {e}')
        sys.exit(1)