Update README.md

README.md

@@ -1,13 +1,471 @@
+# main.py
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+import torch
+from PIL import Image, ImageEnhance
+import os
+import random
+import json
+import argparse
+from pathlib import Path
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from accelerate import Accelerator
+from diffusers import UNet2DConditionModel, AutoencoderKL
+from diffusers.training_utils import EMAModel
+from diffusers import LoraLoaderMixin
+from diffusers.optimization import get_scheduler
+from torch.optim import AdamW
+import math
+import numpy as np
+from huggingface_hub import create_repo, upload_folder
+from huggingface_hub import HfFolder
+from tqdm.auto import tqdm
+
+# --- Configuration ---
+MODEL_NAME = "photo-fluxXL"
+BASE_MODEL = "kudzueye/Boreal"  # Choose one of your base models
+IMAGE_FOLDER = "/content/drive/MyDrive/training_data"  # Replace with your image folder path
+OUTPUT_DIR = "/content/drive/MyDrive/my_lora_models" # Replace with your output folder path
+TRAIN_BATCH_SIZE = 1
+GRADIENT_ACCUMULATION_STEPS = 4
+LEARNING_RATE = 1e-4
+NUM_EPOCHS = 10
+SAVE_STEPS = 500
+SEED = 42
+PUSH_TO_HUB = False # Set to True if you want to push to Hugging Face Hub
+HUB_REPO_ID = "your-username/your-repo-name" # Replace with your Hugging Face repo ID
+GENERATE_AFTER_TRAINING = True # Set to True to generate images after training
+PROMPTS_FILE = "/content/drive/MyDrive/prompts.json" # Replace with your prompts file path
+BATCH_SIZE_GENERATE = 4 # Batch size for generation
+
+# --- Load Base Model ---
+pipe = StableDiffusionPipeline.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=torch.float16,
+    safety_checker=None,
+    requires_safety_checker=False,
+    variant="fp16",
+    use_safetensors=True
+).to("cuda")
+
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(
+    pipe.scheduler.config,
+    algorithm_type="dpmsolver++",
+    solver_order=2
+)
+
+pipe.enable_attention_slicing()
+pipe.enable_xformers_memory_efficient_attention()
+
+# --- Prepare Dataset ---
+class FluxDataset(Dataset):
+    def __init__(self, image_folder, transform=None):
+        self.image_paths = {}
+        for category in os.listdir(image_folder):
+            category_path = os.path.join(image_folder, category)
+            if os.path.isdir(category_path):
+                self.image_paths[category] = []
+                for subcategory in os.listdir(category_path):
+                    subcategory_path = os.path.join(category_path, subcategory)
+                    if os.path.isdir(subcategory_path):
+                        self.image_paths[category].extend([os.path.join(subcategory_path, f) for f in os.listdir(subcategory_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))])
+                    elif subcategory.lower().endswith(('.png', '.jpg', '.jpeg')):
+                        self.image_paths[category].append(subcategory_path)
+
+        self.transform = transform
+        self.categories = list(self.image_paths.keys())
+
+    def __len__(self):
+        return max(len(paths) for paths in self.image_paths.values())
+
+    def __getitem__(self, idx):
+        item = {}
+        for category in self.categories:
+            if self.image_paths[category]:
+                image_path = self.image_paths[category][idx % len(self.image_paths[category])]
+                image = Image.open(image_path).convert("RGB")
+                if self.transform:
+                    image = self.transform(image)
+                item[category] = image
+        return item
+
+transform = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+])
+
+dataset = FluxDataset(IMAGE_FOLDER, transform=transform)
+
+# --- Data Loader ---
+dataloader = DataLoader(dataset, batch_size=TRAIN_BATCH_SIZE, shuffle=True)
+
+# --- Prepare LoRA ---
+unet = pipe.unet
+vae = pipe.vae
+text_encoder = pipe.text_encoder
+text_encoder_2 = pipe.text_encoder_2
+
+unet_lora_layers = LoraLoaderMixin.get_lora_layers(unet)
+vae_lora_layers = LoraLoaderMixin.get_lora_layers(vae)
+text_encoder_lora_layers = LoraLoaderMixin.get_lora_layers(text_encoder)
+text_encoder_2_lora_layers = LoraLoaderMixin.get_lora_layers(text_encoder_2)
+
+# --- Optimizer ---
+optimizer = AdamW(
+    [
+        {"params": unet_lora_layers.parameters(), "lr": LEARNING_RATE},
+        {"params": vae_lora_layers.parameters(), "lr": LEARNING_RATE},
+        {"params": text_encoder_lora_layers.parameters(), "lr": LEARNING_RATE},
+        {"params": text_encoder_2_lora_layers.parameters(), "lr": LEARNING_RATE},
+    ]
+)
+
+# --- Scheduler ---
+lr_scheduler = get_scheduler(
+    "cosine",
+    optimizer=optimizer,
+    num_warmup_steps=math.ceil(len(dataloader) * NUM_EPOCHS * 0.1),
+    num_training_steps=len(dataloader) * NUM_EPOCHS,
+)
+
+# --- Accelerator ---
+accelerator = Accelerator(gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS, mixed_precision="fp16")
+
+unet, vae, text_encoder, text_encoder_2, optimizer, dataloader, lr_scheduler = accelerator.prepare(
+    unet, vae, text_encoder, text_encoder_2, optimizer, dataloader, lr_scheduler
+)
+
+# --- Training Loop ---
+progress_bar = tqdm(range(len(dataloader) * NUM_EPOCHS), desc="Training")
+global_step = 0
+
+for epoch in range(NUM_EPOCHS):
+    for batch in dataloader:
+        with accelerator.accumulate(unet, vae, text_encoder, text_encoder_2):
+            latents = vae.encode(batch["body"].to(accelerator.device)).latent_dist.sample()
+            noise = torch.randn_like(latents)
+            timesteps = torch.randint(0, pipe.scheduler.num_train_timesteps, (latents.shape[0],), device=accelerator.device)
+            noisy_latents = pipe.scheduler.add_noise(latents, noise, timesteps)
+
+            # --- Generate Prompt ---
+            prompt_parts = []
+            if "body" in batch:
+                if "belly" in batch:
+                    prompt_parts.append("belly visible")
+                if "body_shape_1" in batch:
+                    prompt_parts.append("body shape 1")
+                if "body_shape_2" in batch:
+                    prompt_parts.append("body shape 2")
+                if "body_type" in batch:
+                    prompt_parts.append("body type")
+                if "body measurements-proportion" in batch:
+                    prompt_parts.append("body measurements-proportion")
+            if "details" in batch:
+                if "eyebrows" in batch:
+                    prompt_parts.append("eyebrows")
+                if "eyelashes_1" in batch:
+                    prompt_parts.append("eyelashes 1")
+                if "eyelashes_2" in batch:
+                    prompt_parts.append("eyelashes 2")
+                if "hair" in batch:
+                    prompt_parts.append("hair")
+                if "lips" in batch:
+                    prompt_parts.append("lips")
+            if "face" in batch:
+                for i in range(1, 18):
+                    if f"face_{i}" in batch:
+                        prompt_parts.append(f"face {i}")
+            if "pose" in batch:
+                for i in range(1, 4):
+                    if f"pose_{i}" in batch:
+                        prompt_parts.append(f"pose {i}")
+            if "skin" in batch:
+                if "skin_tone" in batch:
+                    prompt_parts.append("skin tone")
+            if "textures" in batch:
+                for texture in batch["textures"]:
+                    prompt_parts.append(f"texture {texture}")
+
+            prompt = "a photo of a woman, " + ", ".join(prompt_parts)
+            prompt_embeds = pipe.text_encoder(pipe.tokenizer(prompt, return_tensors="pt").input_ids.to(accelerator.device))[0]
+            prompt_embeds_2 = pipe.text_encoder_2(pipe.tokenizer(prompt, return_tensors="pt").input_ids.to(accelerator.device))[0]
+
+            model_pred = unet(noisy_latents, timesteps, prompt_embeds, prompt_embeds_2).sample
+
+            loss = torch.nn.functional.mse_loss(model_pred, noise)
+            accelerator.backward(loss)
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        progress_bar.update(1)
+        global_step += 1
+
+        if global_step % SAVE_STEPS == 0:
+            if accelerator.is_main_process:
+                print(f"Saving checkpoint at step {global_step}")
+                save_path = os.path.join(OUTPUT_DIR, f"{MODEL_NAME}_{global_step}")
+                accelerator.save_state(save_path)
+                # Save LoRA weights
+                unet_lora_layers.save_pretrained(os.path.join(save_path, "unet_lora"))
+                vae_lora_layers.save_pretrained(os.path.join(save_path, "vae_lora"))
+                text_encoder_lora_layers.save_pretrained(os.path.join(save_path, "text_encoder_lora"))
+                text_encoder_2_lora_layers.save_pretrained(os.path.join(save_path, "text_encoder_2_lora"))
+
+                # Save model card
+                model_card = f"""---
+license: apache-2.0
+language:
+- en
+base_model:
+- {BASE_MODEL}
+pipeline_tag: text-to-image
 ---
+
+# Model Description: {MODEL_NAME}
+
+This LoRa model enhances text-to-image generation with a hyperrealistic style focusing on a specific subject.
+
+## Subject Description:
+(Add detailed subject description here)
+
+## Hyperrealistic Style: {True}
+
+## Base Models:
+This model was trained using the following base model:
+- {BASE_MODEL}
+
+
+## Usage Instructions:
+(Add detailed instructions on how to use this LoRa model here. Include example prompts)
+
+
+## Training Data:
+(Add information on the training data here)
+
+
+## Limitations:
+(List any known limitations of the model)
+
+## Bias and Fairness Considerations:
+(Address potential bias in the model)
+
+## Known Issues:
+(List any known issues)
+
+"""
+                with open(os.path.join(save_path, "model_card.txt"), "w") as f:
+                    f.write(model_card)
+
+# --- Save Final Model ---
+if accelerator.is_main_process:
+    print("Saving final model")
+    save_path = os.path.join(OUTPUT_DIR, f"{MODEL_NAME}_final")
+    accelerator.save_state(save_path)
+    # Save LoRA weights
+    unet_lora_layers.save_pretrained(os.path.join(save_path, "unet_lora"))
+    vae_lora_layers.save_pretrained(os.path.join(save_path, "vae_lora"))
+    text_encoder_lora_layers.save_pretrained(os.path.join(save_path, "text_encoder_lora"))
+    text_encoder_2_lora_layers.save_pretrained(os.path.join(save_path, "text_encoder_2_lora"))
+
+    # Save model card
+    model_card = f"""---
 license: apache-2.0
+language:
+- en
 base_model:
-- Shakker-Labs/FLUX.1-dev-LoRA-add-details
-- prithivMLmods/Flux-Realism-FineDetailed
-- prithivMLmods/Canopus-LoRA-Flux-UltraRealism-2.0
-- kudzueye/boreal-flux-dev-v2
-- kudzueye/Boreal
-- adirik/flux-cinestill
-- Schaffsch/ostris_flux-dev-lora-trainer
-- Jovie/Midjourney
+- {BASE_MODEL}
 pipeline_tag: text-to-image
----
+---
+
+# Model Description: {MODEL_NAME}
+
+This LoRa model enhances text-to-image generation with a hyperrealistic style focusing on a specific subject.
+
+## Subject Description:
+(Add detailed subject description here)
+
+## Hyperrealistic Style: {True}
+
+## Base Models:
+This model was trained using the following base model:
+- {BASE_MODEL}
+
+
+## Usage Instructions:
+(Add detailed instructions on how to use this LoRa model here. Include example prompts)
+
+
+## Training Data:
+(Add information on the training data here)
+
+
+## Limitations:
+(List any known limitations of the model)
+
+## Bias and Fairness Considerations:
+(Address potential bias in the model)
+
+## Known Issues:
+(List any known issues)
+
+"""
+    with open(os.path.join(save_path, "model_card.txt"), "w") as f:
+        f.write(model_card)
+
+# --- Push to Hub ---
+if PUSH_TO_HUB and accelerator.is_main_process:
+    print("Pushing to Hugging Face Hub")
+    repo_id = HUB_REPO_ID
+    repo_url = create_repo(repo_id, exist_ok=True, repo_type="model", token=HfFolder.get_token()).clone_url
+    upload_folder(repo_id=repo_id, folder_path=save_path, token=HfFolder.get_token())
+    print(f"Model pushed to {repo_url}")
+
+# --- FluxLoraModel Class ---
+class FluxLoraModel:
+    def __init__(self, model_path="your-model-path", device="cuda"):
+        self.device = device
+        self.model = StableDiffusionPipeline.from_pretrained(
+            model_path,
+            torch_dtype=torch.float16,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to(device)
+        
+        self.model.scheduler = DPMSolverMultistepScheduler.from_config(
+            self.model.scheduler.config,
+            algorithm_type="dpmsolver++",
+            solver_order=2
+        )
+        
+        self.model.enable_attention_slicing()
+        self.model.enable_xformers_memory_efficient_attention()
+
+        self.quality_modifiers = {
+            'realism': ["hyperrealistic", "photorealistic", "ultra realistic", "octane render", "raw photo", "unedited", "photographic", "35mm film"],
+            'resolution': ["4K UHD", "8K resolution", "ultra high definition", "extremely detailed", "high resolution"],
+            'detail_level': ["ultra detailed", "fine details", "intricate details", "sharp focus", "highly detailed", "maximum detail"]
+        }
+
+        self.texture_modifiers = {
+            'skin_details': ["detailed skin texture", "natural skin pores", "realistic skin subsurface scattering", "fine skin details"],
+            'clothing_details': ["detailed fabric texture", "intricate fabric weave", "realistic cloth folds", "natural fabric wrinkles"],
+            'hair_details': {
+                'general_quality': ["ultra detailed hair strands", "photorealistic hair texture", "volumetric hair rendering"],
+                'hair_types': {
+                    'straight': ["silky straight hair", "smooth hair texture"],
+                    'wavy': ["natural wave pattern", "defined hair waves"],
+                    'curly': ["detailed curl pattern", "natural curl definition"],
+                    'coily': ["detailed coil pattern", "natural coil definition"]
+                }
+            },
+            'eye_details': {
+                'general_quality': ["ultra detailed iris", "photorealistic eyes", "8K eye details"],
+                'iris_details': ["detailed iris patterns", "intricate iris fibers"],
+                'eye_properties': {
+                    'reflection': ["natural catchlights", "realistic eye reflections"],
+                    'moisture': ["natural eye moisture", "subtle tear film"],
+                    'depth': ["volumetric eye depth", "realistic eye socket depth"]
+                }
+            }
+        }
+
+    def enhance_prompt(self, base_prompt):
+        realism_mod = ", ".join(random.sample(self.quality_modifiers['realism'], 3))
+        resolution_mod = ", ".join(random.sample(self.quality_modifiers['resolution'], 2))
+        detail_mod = ", ".join(random.sample(self.quality_modifiers['detail_level'], 3))
+        
+        enhanced_prompt = f"{base_prompt}, {realism_mod}, {resolution_mod}, {detail_mod}, masterpiece, professional photography"
+        
+        if any(word in base_prompt.lower() for word in ["person", "portrait", "face"]):
+            skin_mod = ", ".join(random.sample(self.texture_modifiers['skin_details'], 2))
+            eye_mod = ", ".join(random.sample(self.texture_modifiers['eye_details']['general_quality'], 2))
+            enhanced_prompt = f"{enhanced_prompt}, {skin_mod}, {eye_mod}"
+
+        return enhanced_prompt
+
+    def generate_image(self, prompt, negative_prompt="", num_images=1, steps=50, cfg_scale=8.5, width=2048, height=2048, seed=None, output_dir="outputs"):
+        default_negative = "blur, haze, soft, deformed, low quality, low resolution, noise, grainy, bad details"
+        enhanced_negative_prompt = f"{negative_prompt}, {default_negative}"
+        enhanced_prompt = self.enhance_prompt(prompt)
+
+        if width >= 1024 or height >= 1024:
+            self.model.enable_vae_tiling()
+
+        generator = torch.Generator(device=self.device).manual_seed(seed) if seed else None
+
+        images = self.model(
+            prompt=enhanced_prompt,
+            negative_prompt=enhanced_negative_prompt,
+            num_images_per_prompt=num_images,
+            num_inference_steps=steps,
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator
+        ).images
+
+        processed_images = []
+        for img in images:
+            img = img.filter(ImageEnhance.Sharpness(1.2))
+            img = img.filter(ImageEnhance.Contrast(1.1))
+            processed_images.append(img)
+
+        os.makedirs(output_dir, exist_ok=True)
+        saved_paths = []
+        for i, image in enumerate(processed_images):
+            path = os.path.join(output_dir, f"flux_4k_detailed_{i}.png")
+            image.save(path, "PNG", quality=100, optimize=True)
+            saved_paths.append(path)
+            
+        return processed_images, saved_paths
+
+    def generate_4k_portrait(self, prompt, **kwargs):
+        return self.generate_image(prompt=prompt, width=3840, height=2160, steps=60, cfg_scale=9.0, **kwargs)
+
+    @staticmethod
+    def image_grid(imgs, rows, cols):
+        w, h = imgs[0].size
+        grid = Image.new('RGB', size=(cols * w, rows * h))
+        for i, img in enumerate(imgs):
+            grid.paste(img, box=(i % cols * w, i // cols * h))
+        return grid
+
+def batch_generate(prompts_file, output_dir="batch_outputs", batch_size=4, **kwargs):
+    model = FluxLoraModel()
+    with open(prompts_file, 'r') as f:
+        prompts = json.load(f)
+    
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    for i, prompt in enumerate(prompts):
+        try:
+            images, paths = model.generate_image(
+                prompt=prompt,
+                num_images=batch_size,
+                output_dir=str(output_dir / f"prompt_{i}"),
+                **kwargs
+            )
+            grid = model.image_grid(images, rows=batch_size//2, cols=2)
+            grid.save(output_dir / f"prompt_{i}_grid.png")
+        except Exception as e:
+            print(f"Error processing prompt {i}: {str(e)}")
+
+# --- Generation after training ---
+if GENERATE_AFTER_TRAINING and accelerator.is_main_process:
+    print("Generating images after training...")
+    # Load the trained LoRA model
+    model = FluxLoraModel(model_path=os.path.join(OUTPUT_DIR, f"{MODEL_NAME}_final"))
+    batch_generate(PROMPTS_FILE, output_dir="generated_images", batch_size=BATCH_SIZE_GENERATE)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prompts", type=str, required=False)
+    parser.add_argument("--output", type=str, default="outputs")
+    parser.add_argument("--batch-size", type=int, default=4)
+    args = parser.parse_args()
+    
+    if args.prompts:
+        batch_generate(args.prompts, args.output, args.batch_size)