| |
| """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" |
|
|
| |
| 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" |
|
|
| |
| SHOTS = [ |
| |
| ("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"), |
|
|
| |
| ("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 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 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"), |
|
|
| |
| ("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() |
|
|
