forked from DreamLab-AI/origin-logseq-AR
-
Notifications
You must be signed in to change notification settings - Fork 19
Expand file tree
/
Copy pathgenerate_cats.py
More file actions
executable file
·265 lines (222 loc) · 8.11 KB
/
generate_cats.py
File metadata and controls
executable file
·265 lines (222 loc) · 8.11 KB
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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
#!/usr/bin/env python3
"""Generate 4 highly differentiated cat images using FLUX model in ComfyUI."""
import requests
import json
import time
import base64
import random
from pathlib import Path
COMFYUI_URL = "http://localhost:8188"
OUTPUT_DIR = Path("/home/devuser/workspace/project/output/cats")
# Four highly differentiated cat prompts
CAT_PROMPTS = [
{
"name": "cyberpunk_neon_cat",
"prompt": "a futuristic cyberpunk cat with glowing neon cyan and magenta eyes, circuit patterns on fur, holographic collar, metallic whiskers, standing in a neon-lit rain-soaked alley, ultra detailed, 8k, cinematic lighting, dramatic shadows",
"width": 1024,
"height": 1024,
},
{
"name": "watercolor_garden_cat",
"prompt": "a soft watercolor painting of an elegant Persian cat sitting in a blooming spring garden, delicate brushstrokes, pastel pink and purple flowers, dreamy atmosphere, artistic, flowing colors, gentle lighting, impressionist style",
"width": 1024,
"height": 1024,
},
{
"name": "glass_sculpture_cat",
"prompt": "a translucent 3D rendered glass sculpture of a sitting cat, crystal clear material with rainbow light refraction, smooth curves, reflective surfaces, studio lighting, white background, photorealistic rendering, elegant and minimalist",
"width": 1024,
"height": 1024,
},
{
"name": "noir_detective_cat",
"prompt": "a vintage film noir style black and white photograph of a sophisticated cat detective wearing a fedora hat and trench coat, sitting at a desk with dramatic venetian blind shadows, cigarette smoke, 1940s atmosphere, high contrast, moody lighting",
"width": 1024,
"height": 1024,
}
]
def create_flux_workflow(prompt, width=1024, height=1024, seed=None):
"""Create a FLUX workflow JSON for ComfyUI."""
if seed is None:
seed = random.randint(0, 2**32 - 1)
workflow = {
"6": {
"inputs": {
"text": prompt,
"clip": ["30", 1]
},
"class_type": "CLIPTextEncode",
"_meta": {"title": "CLIP Text Encode (Positive)"}
},
"8": {
"inputs": {
"samples": ["31", 0],
"vae": ["30", 2]
},
"class_type": "VAEDecode"
},
"9": {
"inputs": {
"filename_prefix": "ComfyUI",
"images": ["8", 0]
},
"class_type": "SaveImage"
},
"27": {
"inputs": {
"width": width,
"height": height,
"batch_size": 1
},
"class_type": "EmptySD3LatentImage"
},
"30": {
"inputs": {
"ckpt_name": "flux1-schnell-fp8.safetensors"
},
"class_type": "CheckpointLoaderSimple"
},
"31": {
"inputs": {
"seed": seed,
"steps": 4, # Schnell is optimized for 4 steps
"cfg": 1.0,
"sampler_name": "euler",
"scheduler": "simple",
"denoise": 1.0,
"model": ["30", 0],
"positive": ["35", 0],
"negative": ["33", 0],
"latent_image": ["27", 0]
},
"class_type": "KSampler"
},
"33": {
"inputs": {
"text": "blurry, low quality, distorted, ugly, bad anatomy",
"clip": ["30", 1]
},
"class_type": "CLIPTextEncode",
"_meta": {"title": "CLIP Text Encode (Negative)"}
},
"35": {
"inputs": {
"guidance": 3.5,
"conditioning": ["6", 0]
},
"class_type": "FluxGuidance"
}
}
return workflow
def queue_prompt(workflow):
"""Queue a prompt to ComfyUI."""
payload = {"prompt": workflow}
response = requests.post(f"{COMFYUI_URL}/prompt", json=payload)
return response.json()
def get_history(prompt_id):
"""Get the history for a prompt ID."""
response = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
return response.json()
def get_image(filename, subfolder, folder_type):
"""Download an image from ComfyUI."""
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
response = requests.get(f"{COMFYUI_URL}/view", params=data)
return response.content
def wait_for_completion(prompt_id, timeout=600, poll_interval=2):
"""Wait for a prompt to complete."""
start_time = time.time()
while time.time() - start_time < timeout:
history = get_history(prompt_id)
if prompt_id in history:
prompt_data = history[prompt_id]
if "outputs" in prompt_data:
return prompt_data
time.sleep(poll_interval)
print(".", end="", flush=True)
raise TimeoutError(f"Prompt {prompt_id} did not complete within {timeout}s")
def generate_cat_image(cat_config):
"""Generate a single cat image."""
print(f"\n{'='*60}")
print(f"Generating: {cat_config['name']}")
print(f"Prompt: {cat_config['prompt'][:80]}...")
print(f"{'='*60}")
# Create workflow
workflow = create_flux_workflow(
cat_config["prompt"],
cat_config["width"],
cat_config["height"]
)
# Queue the prompt
print("Queueing prompt...")
result = queue_prompt(workflow)
prompt_id = result.get("prompt_id")
if not prompt_id:
print(f"Error: {result}")
return None
print(f"Prompt ID: {prompt_id}")
print("Waiting for completion", end="")
# Wait for completion
try:
prompt_data = wait_for_completion(prompt_id, timeout=600)
print("\nGeneration complete!")
except TimeoutError as e:
print(f"\nError: {e}")
return None
# Get the image
outputs = prompt_data.get("outputs", {})
for node_id, node_output in outputs.items():
if "images" in node_output:
for image_data in node_output["images"]:
filename = image_data["filename"]
subfolder = image_data.get("subfolder", "")
folder_type = image_data.get("type", "output")
print(f"Downloading: {filename}")
image_bytes = get_image(filename, subfolder, folder_type)
# Save the image
output_path = OUTPUT_DIR / f"{cat_config['name']}.png"
output_path.write_bytes(image_bytes)
print(f"Saved to: {output_path}")
return str(output_path)
print("No image found in output")
return None
def main():
"""Generate all 4 cat images."""
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
print("\n" + "="*60)
print("FLUX CAT IMAGE GENERATOR")
print("Generating 4 highly differentiated cat images")
print("="*60)
# Check if ComfyUI is running
try:
response = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
print(f"\n✓ ComfyUI is running (v{response.json()['system']['comfyui_version']})")
except Exception as e:
print(f"\n✗ Error: ComfyUI is not responding at {COMFYUI_URL}")
print(f" Details: {e}")
return
# Generate each cat image
results = []
for i, cat_config in enumerate(CAT_PROMPTS, 1):
print(f"\n[{i}/{len(CAT_PROMPTS)}]", end=" ")
output_path = generate_cat_image(cat_config)
results.append({
"name": cat_config["name"],
"path": output_path,
"prompt": cat_config["prompt"]
})
# Brief pause between generations
if i < len(CAT_PROMPTS):
time.sleep(2)
# Summary
print("\n" + "="*60)
print("GENERATION COMPLETE")
print("="*60)
for i, result in enumerate(results, 1):
status = "✓" if result["path"] else "✗"
print(f"{status} [{i}] {result['name']}")
if result["path"]:
print(f" Path: {result['path']}")
print(f"\nAll images saved to: {OUTPUT_DIR}")
print("="*60)
if __name__ == "__main__":
main()