Update README.md

README.md

@@ -1,471 +1,15 @@
-# 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
+- es
+- fr
 base_model:
-- {BASE_MODEL}
+- kudzueye/Boreal
+- adirik/flux-cinestill
+- Shakker-Labs/FLUX.1-dev-LoRA-add-details
+- prithivMLmods/Flux-Realism-FineDetailed
+- prithivMLmods/Canopus-LoRA-Flux-UltraRealism-2.0
+- Jovie/Midjourney
 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)
+---