run_consistentId_sdxl.py

import random
import shutil
from gradio_client import Client, handle_file
from tqdm import tqdm
import re
import os
import pathlib
from datasets import Dataset, Image as HuggingFaceImage
from PIL import Image

prefix_quick_list = [
    "detailed photo of",
    "amateur photo of",
    "flicker 2008 photo of",
    "fantastic artwork of",
    "vintage photograph of",
    "Unreal 5 render of",
    "surrealist painting of",
    "professional advertising design of"
]

subject_quick_list = [
    "a handsome man img"
]

illumination_environment_quick_list = [
    "sunshine from window", "neon light, city", "sunset over sea", "golden time", "sci-fi RGB glowing, cyberpunk", "natural lighting",
    "warm atmosphere, at home, bedroom", "magic lit", "evil, gothic, in a cave", "light and shadow", "shadow from window", "soft studio lighting",
    "home atmosphere, cozy bedroom illumination", "neon, Wong Kar-wai, warm", "moonlight through curtains", "stormy sky lighting",
    "underwater glow, deep sea", "foggy forest at dawn", "golden hour in a meadow", "rainbow reflections, neon", "cozy candlelight",
    "apocalyptic, smoky atmosphere", "red glow, emergency lights", "mystical glow, enchanted forest", "campfire light", "harsh, industrial lighting",
    "sunrise in the mountains", "evening glow in the desert", "moonlight in a dark alley", "golden glow at a fairground",
    "midnight in the forest", "purple and pink hues at twilight", "foggy morning, muted light", "candle-lit room, rustic vibe",
    "fluorescent office lighting", "lightning flash in storm", "night, cozy warm light from fireplace", "ethereal glow, magical forest",
    "dusky evening on a beach", "afternoon light filtering through trees", "blue neon light, urban street", "red and blue police lights in rain",
    "aurora borealis glow, arctic landscape", "sunrise through foggy mountains", "golden hour on a city skyline", "mysterious twilight, heavy mist",
    "early morning rays, forest clearing", "colorful lantern light at festival", "soft glow through stained glass", "harsh spotlight in dark room",
    "mellow evening glow on a lake", "crystal reflections in a cave", "vibrant autumn lighting in a forest", "gentle snowfall at dusk",
    "hazy light of a winter morning", "soft, diffused foggy glow", "underwater luminescence", "rain-soaked reflections in city lights",
    "golden sunlight streaming through trees", "fireflies lighting up a summer night", "glowing embers from a forge", "dim candlelight in a gothic castle",
    "midnight sky with bright starlight", "warm sunset in a rural village", "flickering light in a haunted house", "desert sunset with mirage-like glow",
    "golden beams piercing through storm clouds"
]

background_quick_list = [
    "cars and people", "a cozy bed and a lamp", "a forest clearing with mist", "a bustling marketplace", "a quiet beach at dusk",
    "an old, cobblestone street", "a futuristic cityscape", "a tranquil lake with mountains", "a mysterious cave entrance",
    "bookshelves and plants in the background", "an ancient temple in ruins", "tall skyscrapers and neon signs", "a starry sky over a desert",
    "a bustling café", "rolling hills and farmland", "a modern living room with a fireplace", "an abandoned warehouse",
    "a picturesque mountain range", "a starry night sky", "the interior of a futuristic spaceship", "the cluttered workshop of an inventor",
    "the glowing embers of a bonfire", "a misty lake surrounded by trees", "an ornate palace hall", "a busy street market", "a vast desert landscape",
    "a peaceful library corner", "bustling train station", "a mystical, enchanted forest", "an underwater reef with colorful fish", "a quiet rural village",
    "a sandy beach with palm trees", "a vibrant coral reef, teeming with life", "snow-capped mountains in distance", "a stormy ocean, waves crashing",
    "a rustic barn in open fields", "a futuristic lab with glowing screens", "a dark, abandoned castle", "the ruins of an ancient civilization",
    "a bustling urban street in rain", "an elegant grand ballroom", "a sprawling field of wildflowers", "a dense jungle with sunlight filtering through",
    "a dimly lit, vintage bar", "an ice cave with sparkling crystals", "a serene riverbank at sunset", "a narrow alley with graffiti walls",
    "a peaceful zen garden with koi pond", "a high-tech control room", "a quiet mountain village at dawn", "a lighthouse on a rocky coast",
    "a rainy street with flickering lights", "a frozen lake with ice formations", "an abandoned theme park", "a small fishing village on a pier",
    "rolling sand dunes in a desert", "a dense forest with towering redwoods", "a snowy cabin in the mountains", "a mystical cave with bioluminescent plants",
    "a castle courtyard under moonlight", "a bustling open-air night market", "an old train station with steam", "a tranquil waterfall surrounded by trees",
    "a vineyard in the countryside", "a quaint medieval village", "a bustling harbor with boats", "a high-tech futuristic mall",
    "a lush tropical rainforest"
]

style_list = ['(No style)', 'Cinematic', 'Disney Charactor', 'Digital Art', 'Photographic (Default)', 'Fantasy art', 'Neonpunk', 'Enhance', 'Comic book', 'Lowpoly', 'Line art']

def clean_filename(prompt):
    # Replace spaces with underscores
    filename = prompt.replace(' ', '_')
    # Remove or replace non-alphanumeric characters (except underscores)
    filename = re.sub(r'[^a-zA-Z0-9_]', '', filename)
    return filename + '.png'

output_dir = 'phm2_output_images'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

def generate_unique_random_string(used_prompts):
    while True:
        prefix = random.choice(prefix_quick_list)
        subject = random.choice(subject_quick_list)
        illumination = random.choice(illumination_environment_quick_list)
        background = random.choice(background_quick_list)
        
        prompt = f"{prefix} {subject}. {illumination}. {background} in the background."
        if prompt not in used_prompts:
            used_prompts.add(prompt)
            return {
                "prompt": prompt,
                "prefix": prefix,
                "subject": subject,
                "illumination": illumination,
                "background": background
            }

num_images = 10
used_prompts = set()

dataset_dict = {
    "image": [],
    "prefix": [],
    "subject": [],
    "illumination": [],
    "background": [],
    "prompt": [],
    #"style_name": []
}

template_image_path = '/home/featurize/ConsistentID-SDXL/xiangxiang_IMG.jpeg'
template_image = handle_file(template_image_path)

client = Client("http://127.0.0.1:7860")

for _ in tqdm(range(num_images), desc="Processing prompts"):
    prompt_details = generate_unique_random_string(used_prompts)
    style_name = random.choice(style_list)
    
    result = client.predict(
        selected_template_images=template_image_path,
        costum_image=template_image,
        prompt=prompt_details["prompt"],
        negative_prompt="nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry",
        prompt_selected='A young, handsome boy in a wedding suit, with a slender figure and a cute face',
        retouching=False,
        width=864,
        height=1152,
        merge_steps=30,
        seed_set=2024,
        api_name="/process"
    )
    
    # Save image and update dataset_dict
    image_path = result
    
    image = Image.open(image_path)
    
    dataset_dict["image"].append(image)
    dataset_dict["prefix"].append(prompt_details["prefix"])
    dataset_dict["subject"].append(prompt_details["subject"])
    dataset_dict["illumination"].append(prompt_details["illumination"])
    dataset_dict["background"].append(prompt_details["background"])
    dataset_dict["prompt"].append(prompt_details["prompt"])
    #dataset_dict["style_name"].append(style_name)
    
dataset = Dataset.from_dict(dataset_dict)
dataset.save_to_disk("csdxl_dataset")

from datasets import load_from_disk 
load_from_disk("csdxl_dataset").push_to_hub("svjack/Prince_Xiang_ConsistentID_SDXL_10")