File size: 10,554 Bytes
a70eb3d | 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 | #!/usr/bin/env python3
"""Generate diverse Indian woman character dataset using Flux 2 Dev for LoRA training."""
import json
import urllib.request
import time
import os
COMFYUI_URL = "http://127.0.0.1:80"
DATASET_DIR = "/home/azureuser/ai-toolkit/character_dataset"
TRIGGER = "ohwx"
# Consistent identity description used in EVERY prompt
IDENTITY = f"a young 21 year old Indian woman called {TRIGGER}, long straight black hair, brown eyes, light brown skin, soft facial features, natural beauty"
# (prompt_suffix, filename, caption)
SHOTS = [
# === ANGLES / HEAD POSITIONS ===
("front facing portrait, looking directly at camera, neutral expression, white background, studio lighting, photorealistic", "angle_front", "front facing portrait, looking at camera, neutral expression, studio lighting"),
("side profile portrait, left side of face visible, looking left, clean background, studio lighting, photorealistic", "angle_left_profile", "side profile, left side visible, looking left, studio lighting"),
("side profile portrait, right side of face visible, looking right, clean background, studio lighting, photorealistic", "angle_right_profile", "side profile, right side visible, looking right, studio lighting"),
("three quarter view portrait, slightly turned to the left, soft lighting, clean background, photorealistic", "angle_34_left", "three quarter view, slightly turned left, soft lighting"),
("three quarter view portrait, slightly turned to the right, soft lighting, clean background, photorealistic", "angle_34_right", "three quarter view, slightly turned right, soft lighting"),
("looking up, chin slightly raised, soft lighting from above, clean background, photorealistic portrait", "angle_looking_up", "looking up, chin raised, soft overhead lighting"),
("looking down, reading a book, natural indoor lighting, soft focus background, candid photo", "angle_looking_down", "looking down reading a book, natural indoor lighting, candid"),
("back view, looking over shoulder at camera, long black hair visible, outdoor setting, natural light", "angle_back", "back view looking over shoulder, long black hair visible, outdoor natural light"),
("back of head view, long straight black hair falling down, standing outdoors, natural daylight", "angle_back_full", "back of head, long straight black hair, standing outdoors"),
("tilted head portrait, slight head tilt to the right, gentle smile, soft lighting, photorealistic", "angle_tilt", "tilted head portrait, slight tilt right, gentle smile"),
# === EMOTIONS ===
("smiling brightly, genuine happy smile showing teeth, warm expression, natural light, photorealistic portrait", "emo_happy", "smiling brightly, genuine happy smile, warm expression, natural light"),
("sad expression, looking down slightly, melancholic mood, soft moody lighting, photorealistic portrait", "emo_sad", "sad expression, looking down, melancholic, soft moody lighting"),
("laughing candidly, eyes slightly closed, natural joyful moment, outdoor sunlight, candid photo", "emo_laugh", "laughing candidly, joyful moment, outdoor sunlight, candid"),
("serious contemplative expression, deep in thought, looking into distance, dramatic lighting, photorealistic", "emo_serious", "serious contemplative expression, deep in thought, dramatic lighting"),
("surprised expression, eyes wide, mouth slightly open, bright lighting, photorealistic portrait", "emo_surprised", "surprised expression, eyes wide, bright lighting"),
("shy smile, looking slightly away from camera, subtle smile, soft warm lighting, photorealistic", "emo_shy", "shy smile, looking slightly away, subtle smile, soft warm lighting"),
("confident expression, strong gaze at camera, slight smirk, professional lighting, photorealistic", "emo_confident", "confident expression, strong gaze, slight smirk, professional lighting"),
("peaceful serene expression, eyes closed, meditating, soft natural light, photorealistic portrait", "emo_peaceful", "peaceful serene expression, eyes closed, meditating, soft light"),
# === HALF BODY / DIFFERENT OUTFITS ===
("half body photo, wearing traditional Indian salwar kameez, standing in a garden, natural sunlight, photorealistic", "outfit_salwar", "half body, wearing traditional Indian salwar kameez, garden, sunlight"),
("half body photo, wearing modern casual jeans and white t-shirt, urban street background, daylight, photorealistic", "outfit_casual", "half body, casual jeans and white t-shirt, urban street, daylight"),
("half body photo, wearing formal black blazer and white shirt, office setting, professional lighting, photorealistic", "outfit_formal", "half body, formal black blazer, office setting, professional lighting"),
("half body photo, wearing a red saree with gold border, elegant pose, warm indoor lighting, photorealistic", "outfit_saree", "half body, wearing red saree with gold border, elegant pose, warm indoor light"),
("half body photo, wearing athletic sportswear, gym background, bright fluorescent lighting, photorealistic", "outfit_sport", "half body, athletic sportswear, gym, bright lighting"),
("half body photo, wearing a cozy sweater, sitting by window, rainy day outside, soft natural light, photorealistic", "outfit_cozy", "half body, cozy sweater, sitting by window, rainy day, soft light"),
# === FULL BODY / DIFFERENT SCENES ===
("full body photo, walking on a beach at sunset, wearing summer dress, golden hour light, barefoot on sand, photorealistic", "full_beach", "full body, walking on beach at sunset, summer dress, golden hour"),
("full body photo, standing in a city street, wearing winter jacket, evening city lights, urban photography, photorealistic", "full_city", "full body, city street, winter jacket, evening city lights"),
("full body photo, sitting cross legged on grass in a park, casual outfit, dappled sunlight through trees, photorealistic", "full_park", "full body, sitting on grass in park, casual outfit, dappled sunlight"),
("full body photo, dancing in the rain, white dress, joyful expression, dramatic wet look, photorealistic", "full_rain", "full body, dancing in rain, white dress, joyful, dramatic"),
("full body photo, standing at a temple entrance, traditional Indian outfit, warm morning light, photorealistic", "full_temple", "full body, standing at temple entrance, traditional outfit, morning light"),
# === DIFFERENT LIGHTING ===
("portrait, dramatic chiaroscuro lighting, half face in shadow, artistic, high contrast, photorealistic", "light_dramatic", "portrait, dramatic chiaroscuro lighting, half face in shadow, high contrast"),
("portrait, soft golden hour backlight, hair glowing, warm tones, outdoor, lens flare, photorealistic", "light_golden", "portrait, golden hour backlight, hair glowing, warm tones, outdoor"),
("portrait, bright overcast daylight, flat even lighting, outdoor, clean natural look, photorealistic", "light_overcast", "portrait, bright overcast daylight, even lighting, outdoor, natural"),
("portrait, neon colored lights reflecting on face, nighttime, urban setting, cinematic, photorealistic", "light_neon", "portrait, neon lights on face, nighttime, urban, cinematic"),
("portrait, soft warm candlelight, intimate setting, warm orange tones, photorealistic", "light_candle", "portrait, soft candlelight, intimate setting, warm orange tones"),
]
def queue_prompt(prompt_text, filename, seed):
workflow = {
"1": {"class_type": "UNETLoader", "inputs": {"unet_name": "flux2_dev_fp8mixed.safetensors", "weight_dtype": "default"}},
"2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "mistral_3_small_flux2_bf16.safetensors", "type": "flux2"}},
"3": {"class_type": "VAELoader", "inputs": {"vae_name": "flux2-vae.safetensors"}},
"9": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["2", 0], "text": prompt_text}},
"10": {"class_type": "EmptyFlux2LatentImage", "inputs": {"width": 1024, "height": 1024, "batch_size": 1}},
"11": {"class_type": "KSampler", "inputs": {
"model": ["1", 0], "positive": ["9", 0], "negative": ["9", 0],
"latent_image": ["10", 0], "seed": seed,
"control_after_generate": "fixed", "steps": 20, "cfg": 1.0,
"sampler_name": "euler", "scheduler": "simple", "denoise": 1.0
}},
"12": {"class_type": "VAEDecode", "inputs": {"samples": ["11", 0], "vae": ["3", 0]}},
"13": {"class_type": "SaveImage", "inputs": {"images": ["12", 0], "filename_prefix": f"ds_{filename}"}}
}
data = json.dumps({"prompt": workflow}).encode()
req = urllib.request.Request(f'{COMFYUI_URL}/prompt', data=data, headers={'Content-Type': 'application/json'})
resp = urllib.request.urlopen(req)
return json.loads(resp.read())['prompt_id']
def wait_for_completion(prompt_id, timeout=600):
start = time.time()
while time.time() - start < timeout:
req = urllib.request.Request(f'{COMFYUI_URL}/history/{prompt_id}')
resp = urllib.request.urlopen(req)
history = json.loads(resp.read())
if prompt_id in history:
h = history[prompt_id]
s = h.get('status', {}).get('status_str', '')
if s == 'success':
for out in h['outputs'].values():
if 'images' in out:
return out['images'][0]['filename']
elif s == 'error':
return None
time.sleep(2)
return None
def main():
print(f"Generating {len(SHOTS)} images with Flux 2 Dev...\n")
for i, (suffix, name, caption) in enumerate(SHOTS):
seed = 42424 + i * 997
full_prompt = f"{IDENTITY}, {suffix}"
full_caption = f"photo of {IDENTITY}, {caption}"
print(f"[{i+1}/{len(SHOTS)}] {name}")
prompt_id = queue_prompt(full_prompt, name, seed)
filename = wait_for_completion(prompt_id)
if filename:
src = f"/home/azureuser/ComfyUI/output/{filename}"
dst = os.path.join(DATASET_DIR, f"{name}.png")
os.system(f"cp '{src}' '{dst}'")
with open(os.path.join(DATASET_DIR, f"{name}.txt"), 'w') as f:
f.write(full_caption)
print(f" -> OK")
else:
print(f" -> FAILED")
total = len([f for f in os.listdir(DATASET_DIR) if f.endswith('.png')])
print(f"\nDone! {total} images in {DATASET_DIR}")
if __name__ == "__main__":
main()
|