Update main.py

main.py

@@ -1,62 +1,98 @@
 import os
+import shutil
+import json
 import torch
-from torch import autocast
-from diffusers import StableDiffusionPipeline, DDIMScheduler
-from train_dreambooth import train_dreambooth
-
-class DreamboothApp:
-    def __init__(self, model_path, pretrained_model_name_or_path="runwayml/stable-diffusion-v1-5"):
-        self.model_path = model_path
-        self.pretrained_model_name_or_path = pretrained_model_name_or_path
-        self.pipe = None
-        self.g_cuda = torch.Generator(device='cuda')
-
-    def load_model(self):
-        self.pipe = StableDiffusionPipeline.from_pretrained(self.model_path, safety_checker=None, torch_dtype=torch.float16).to("cuda")
-        self.pipe.scheduler = DDIMScheduler.from_config(self.pipe.scheduler.config)
-        self.pipe.enable_xformers_memory_efficient_attention()
-
-    def train(self, instance_data_dir, class_data_dir, instance_prompt, class_prompt, num_class_images=50, max_train_steps=800, output_dir="stable_diffusion_weights"):
-        concepts_list = [
-            {
-                "instance_prompt": instance_prompt,
-                "class_prompt": class_prompt,
-                "instance_data_dir": instance_data_dir,
-                "class_data_dir": class_data_dir
-            }
-        ]
-        train_dreambooth(pretrained_model_name_or_path=self.pretrained_model_name_or_path,
-                         pretrained_vae_name_or_path="stabilityai/sd-vae-ft-mse",
-                         output_dir=output_dir,
-                         revision="fp16",
-                         with_prior_preservation=True,
-                         prior_loss_weight=1.0,
-                         seed=1337,
-                         resolution=512,
-                         train_batch_size=1,
-                         train_text_encoder=True,
-                         mixed_precision="fp16",
-                         use_8bit_adam=True,
-                         gradient_accumulation_steps=1,
-                         learning_rate=1e-6,
-                         lr_scheduler="constant",
-                         lr_warmup_steps=0,
-                         num_class_images=num_class_images,
-                         sample_batch_size=4,
-                         max_train_steps=max_train_steps,
-                         save_interval=10000,
-                         save_sample_prompt=instance_prompt,
-                         concepts_list=concepts_list)
-        self.model_path = output_dir
-
-    def inference(self, prompt, negative_prompt, num_samples, height=512, width=512, num_inference_steps=50, guidance_scale=7.5, seed=None):
-        if seed is not None:
-            self.g_cuda.manual_seed(seed)
-        with autocast("cuda"), torch.inference_mode():
-            return self.pipe(
-                prompt, height=int(height), width=int(width),
-                negative_prompt=negative_prompt,
-                num_images_per_prompt=int(num_samples),
-                num_inference_steps=int(num_inference_steps), guidance_scale=guidance_scale,
-                generator=self.g_cuda
-            ).images
+import random
+from pathlib import Path
+from torch.utils.data import Dataset
+from torchvision import transforms
+from diffusers import StableDiffusionPipeline, DDIMScheduler, UNet2DConditionModel, AutoencoderKL, DDPMScheduler
+from transformers import CLIPTextModel, CLIPTokenizer
+from accelerate import Accelerator
+from tqdm.auto import tqdm
+from PIL import Image
+
+class CustomDataset(Dataset):
+    def __init__(self, data_dir, prompt, tokenizer, size=512, center_crop=False):
+        self.data_dir = Path(data_dir)
+        self.prompt = prompt
+        self.tokenizer = tokenizer
+        self.size = size
+        self.center_crop = center_crop
+
+        self.image_transforms = transforms.Compose([
+            transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
+            transforms.CenterCrop(size) if center_crop else transforms.RandomCrop(size),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5], [0.5])
+        ])
+
+        self.images = [f for f in self.data_dir.iterdir() if f.is_file() and not str(f).endswith(".txt")]
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx):
+        image_path = self.images[idx]
+        image = Image.open(image_path)
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+
+        image = self.image_transforms(image)
+        prompt_ids = self.tokenizer(
+            self.prompt, padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length
+        ).input_ids
+
+        return {"image": image, "prompt_ids": prompt_ids}
+
+def fine_tune_model(instance_data_dir, instance_prompt, model_name, output_dir, seed=1337, resolution=512, train_batch_size=1, max_train_steps=800):
+    # Setup
+    accelerator = Accelerator(cpu=True)
+    set_seed(seed)
+    
+    tokenizer = CLIPTokenizer.from_pretrained(model_name)
+    text_encoder = CLIPTextModel.from_pretrained(model_name)
+    vae = AutoencoderKL.from_pretrained(model_name)
+    unet = UNet2DConditionModel.from_pretrained(model_name)
+    noise_scheduler = DDPMScheduler.from_pretrained(model_name, subfolder="scheduler")
+
+    dataset = CustomDataset(instance_data_dir, instance_prompt, tokenizer, resolution)
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=train_batch_size, shuffle=True)
+
+    optimizer = torch.optim.AdamW(unet.parameters(), lr=1e-6)
+
+    unet, optimizer, dataloader = accelerator.prepare(unet, optimizer, dataloader)
+    vae.to(accelerator.device)
+    text_encoder.to(accelerator.device)
+
+    global_step = 0
+    for step, batch in tqdm(enumerate(dataloader), total=max_train_steps):
+        latents = vae.encode(batch["image"].to(accelerator.device)).latent_dist.sample() * 0.18215
+        noise = torch.randn_like(latents)
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=latents.device).long()
+        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+        encoder_hidden_states = text_encoder(batch["prompt_ids"].to(accelerator.device))[0]
+
+        model_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+
+        loss = torch.nn.functional.mse_loss(model_pred.float(), noise.float(), reduction="mean")
+        accelerator.backward(loss)
+
+        optimizer.step()
+        optimizer.zero_grad()
+        global_step += 1
+        if global_step >= max_train_steps:
+            break
+
+    # Save model
+    unet = accelerator.unwrap_model(unet)
+    unet.save_pretrained(output_dir)
+    vae.save_pretrained(output_dir)
+    text_encoder.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+
+def set_seed(seed):
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)