Add sd35 T2I/I2I generation scripts

generator/I2I/sd35/edit_sd35_image.py

@@ -0,0 +1,104 @@
+from diffusers import StableDiffusion3Img2ImgPipeline
+from pathlib import Path
+from PIL import Image
+import torch
+import re
+
+INPUT_DIR = Path("/home/yzhou8/sd35_t2i_output") 
+OUT_DIR   = Path("/home/yzhou8/sd35_i2i_output")
+MODEL_DIR = Path("/scratch/yzhou8/stable-diffusion-3.5-medium")
+
+NUM_STEPS       = 40
+GUIDANCE_SCALE  = 4.5
+STRENGTH        = 0.6
+ITERATIONS      = 5       
+
+ALLOWED_VARIANTS = {"traditional", "modern", "general", "national", "common"}
+ALLOWED_COUNTRIES = {"china","korea", "united_states", "kenya", "nigeria", "india"}
+CATEGORY_TO_SUB = {
+    "architecture": {"landmark", "house"},
+    "art": {"painting", "sculpture", "dance"},
+    "event": {"sport", "game", "religious_ritual", "festival", "wedding", "funeral"},
+    "fashion": {"clothing", "makeup", "accessories"},
+    "food": {"staple_food", "main_dish", "snack", "dessert", "beverage"},
+    "landscape": {"city", "countryside", "nature"},
+    "people": {"daily_life", "chef", "celebrity", "model", "athlete", "student", "teacher", "bride_and_groom", "president", "farmer", "soldier", "doctor"}, 
+    "wildlife": {"plant", "animal"}
+}
+
+
+NAME_RE = re.compile(r"^(sd35)_([a-z0-9_]+)_([a-z0-9_]+)_([a-z0-9_]+)_([a-z0-9_]+)\.png$")
+
+def build_prompt(country: str, category: str, subcategory: str, variant: str) -> str:
+    if category == "people" or variant == "general":
+        return f"Change the image to represent {subcategory} in {country}."
+    return f"Change the image to represent {variant} {subcategory} in {country}."
+
+def valid_seed_image(p: Path):
+    m = NAME_RE.match(p.name)
+    if not m:
+        return None
+    model, country, category, subcategory, variant = m.groups()
+
+    if model != "sd35":
+        return None
+    if ALLOWED_COUNTRIES and country not in ALLOWED_COUNTRIES:
+        return None
+    if ALLOWED_VARIANTS and variant not in ALLOWED_VARIANTS:
+        return None
+
+    allowed_subs = CATEGORY_TO_SUB.get(category)
+    if not allowed_subs or subcategory not in allowed_subs:
+        return None
+
+    return model, country, category, subcategory, variant
+
+def ensure_mode(img: Image.Image) -> Image.Image:
+    return img.convert("RGB")
+
+
+def main():
+    OUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    pipe = StableDiffusion3Img2ImgPipeline.from_pretrained(
+        MODEL_DIR,
+        torch_dtype=torch.bfloat16,
+    ).to("cuda")
+
+    for img_path in sorted(INPUT_DIR.glob("*.png")):
+        meta = valid_seed_image(img_path)
+        if not meta:
+            continue 
+        model, country, category, subcategory, variant = meta
+
+        prompt = build_prompt(country, category, subcategory, variant)
+
+        current_image = ensure_mode(Image.open(img_path))
+
+        for i in range(1, ITERATIONS + 1):
+            if category == "people":
+                out_name = f"{model}_{country}_{category}_{subcategory}_{i}.png"
+
+            else:
+                out_name = f"{model}_{country}_{category}_{subcategory}_{variant}_{i}.png"
+            out_path = OUT_DIR / out_name
+
+            if out_path.exists():
+                current_image = ensure_mode(Image.open(out_path))
+                continue
+
+            result = pipe(
+                prompt=prompt,
+                image=current_image,
+                strength=STRENGTH,
+                guidance_scale=GUIDANCE_SCALE,
+                num_inference_steps=NUM_STEPS,
+            ).images[0]
+
+            result.save(out_path)
+            current_image = result
+
+            print(f"Saved: {out_path} | Prompt: {prompt}")
+
+if __name__ == "__main__":
+    main()

generator/T2I/sd35/generate_sd35_image.py

@@ -0,0 +1,67 @@
+from diffusers import StableDiffusion3Pipeline
+import pandas as pd
+import torch
+from pathlib import Path
+from tqdm import tqdm
+
+CSV_PATH = "/home/yzhou8/prompt_india.csv"
+OUT_DIR = "/home/yzhou8/test"
+MODEL_DIR = "/scratch/yzhou8/stable-diffusion-3.5-medium"
+
+pipe = StableDiffusion3Pipeline.from_pretrained(
+    MODEL_DIR,
+    torch_dtype=torch.bfloat16,
+)
+pipe = pipe.to("cuda")
+
+def generate_and_save(prompt: str, out_path: Path, num_inference_steps=40, guidance_scale=4.5):
+    """Run inference and save image"""
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    image = pipe(
+        prompt,
+        num_inference_steps=num_inference_steps,
+        guidance_scale=guidance_scale,
+    ).images[0]
+    image.save(f"{out_path}.png")
+
+
+def main():
+    try:
+        df = pd.read_csv(CSV_PATH) 
+        print('✅ Load CSV file')
+    except FileNotFoundError:
+        print('❌ Error')
+        exit()
+    current_category = ""
+    current_subcategory = ""
+    for index, row in tqdm(df.iterrows(), total=df.shape[0], desc='Generating images from CSV'):
+        if pd.notna(row['Category']):
+            current_category = row['Category']
+        if pd.notna(row['Subcategory']):
+            current_subcategory = row['Subcategory']
+        country = row['Country']
+
+        if pd.isna(country):
+            continue
+
+        country_str = str(country).replace(" ", "_").lower()
+        category_str = str(current_category).replace(" ", "_").lower()
+        subcategory_str = str(current_subcategory).replace(" ", "_").lower() 
+        if country_str == "no_country":
+            country_str = "nocountry"
+
+        for prompt_type in ['Traditional Prompt', 'Modern Prompt', 'General Prompt']:
+            prompt_text = row[prompt_type]
+
+            if isinstance(prompt_text, str) and prompt_text.strip():
+                variant_name = prompt_type.replace(" Prompt", "").lower()
+
+                out_path = Path(OUT_DIR) / f"sd35_{country_str}_{category_str}_{subcategory_str}_{variant_name}"
+
+                print(f"Country: {country}, Category: {current_category}, Sub-category: {current_subcategory}, Variant: {variant_name}")
+                print(f"Prompt: {prompt_text}")
+
+                generate_and_save(prompt_text, out_path)
+
+if __name__ == "__main__":
+    main()