add altdiffusion

app.py

@@ -13,7 +13,7 @@ import zipfile
 import tarfile
 import urllib.parse
 import gc
-from diffusers import StableDiffusionPipeline
+# from diffusers import StableDiffusionPipeline
 from huggingface_hub import snapshot_download
 
 
@@ -34,6 +34,8 @@ if(is_gpu_associated):
     model_v1 = snapshot_download(repo_id="multimodalart/sd-fine-tunable")
     model_v2 = snapshot_download(repo_id="stabilityai/stable-diffusion-2")
     model_v2_512 = snapshot_download(repo_id="stabilityai/stable-diffusion-2-base")
+    model_alt = snapshot_download(repo_id="BAAI/AltDiffusion")
+    model_alt_m9 = snapshot_download(repo_id="BAAI/AltDiffusion-m9")
     safety_checker = snapshot_download(repo_id="multimodalart/sd-sc")
     model_to_load = model_v1
 
@@ -69,6 +71,10 @@ def swap_base_model(selected_model):
             model_to_load = model_v1
         elif(selected_model == "v2-768"):
             model_to_load = model_v2
+        elif(selected_model == "alt"):
+            model_to_load = model_alt
+        elif(selected_model == "alt_m9"):
+            model_to_load = model_alt_m9
         else:
             model_to_load = model_v2_512
 
@@ -288,11 +294,13 @@ def train(*inputs):
 pipe_is_set = False
 def generate(prompt, steps):
     torch.cuda.empty_cache()
-    from diffusers import StableDiffusionPipeline
+    # from diffusers import StableDiffusionPipeline
+    from diffusers import DiffusionPipeline
     global pipe_is_set
     if(not pipe_is_set):
         global pipe
-        pipe = StableDiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)
+        # pipe = StableDiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)
+        pipe = DiffusionPipeline.from_pretrained("./output_model", torch_dtype=torch.float16)
         pipe = pipe.to("cuda")
         pipe_is_set = True
         
@@ -477,7 +485,7 @@ with gr.Blocks(css=css) as demo:
     
     with gr.Row() as what_are_you_training:
         type_of_thing = gr.Dropdown(label="What would you like to train?", choices=["object", "person", "style"], value="object", interactive=True)
-        base_model_to_use = gr.Dropdown(label="Which base model would you like to use?", choices=["v1-5", "v2-512", "v2-768"], value="v1-5", interactive=True)
+        base_model_to_use = gr.Dropdown(label="Which base model would you like to use?", choices=["v1-5", "v2-512", "v2-768", "alt", "alt_m9"], value="alt_m9", interactive=True)
     
     #Very hacky approach to emulate dynamically created Gradio components   
     with gr.Row() as upload_your_concept:

train_dreambooth.py

@@ -17,17 +17,34 @@ from torch.utils.data import Dataset
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import set_seed
-from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from diffusers import AutoencoderKL, DDPMScheduler, DiffusionPipeline, UNet2DConditionModel
 from diffusers.optimization import get_scheduler
 from huggingface_hub import HfFolder, Repository, whoami
 from PIL import Image
 from torchvision import transforms
 from tqdm.auto import tqdm
-from transformers import CLIPTextModel, CLIPTokenizer
+from transformers import AutoTokenizer, PretrainedConfig
 
 
 logger = get_logger(__name__)
 
+def import_model_class_from_model_name_or_path(pretrained_model_name_or_path: str):
+    text_encoder_config = PretrainedConfig.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="text_encoder",
+    )
+    model_class = text_encoder_config.architectures[0]
+
+    if model_class == "CLIPTextModel":
+        from transformers import CLIPTextModel
+
+        return CLIPTextModel
+    elif model_class == "RobertaSeriesModelWithTransformation":
+        from diffusers.pipelines.alt_diffusion.modeling_roberta_series import RobertaSeriesModelWithTransformation
+
+        return RobertaSeriesModelWithTransformation
+    else:
+        raise ValueError(f"{model_class} is not supported.")
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Simple example of a training script.")
@@ -471,7 +488,7 @@ def run_training(args_imported):
 
         if cur_class_images < args.num_class_images:
             torch_dtype = torch.float16 if accelerator.device.type == "cuda" else torch.float32
-            pipeline = StableDiffusionPipeline.from_pretrained(
+            pipeline = DiffusionPipeline.from_pretrained(
                 args.pretrained_model_name_or_path, torch_dtype=torch_dtype
             )
             pipeline.set_progress_bar_config(disable=True)
@@ -517,20 +534,27 @@ def run_training(args_imported):
 
     # Load the tokenizer
     if args.tokenizer_name:
-        tokenizer = CLIPTokenizer.from_pretrained(args.tokenizer_name)
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, use_fast=False)
     elif args.pretrained_model_name_or_path:
-        tokenizer = CLIPTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer")
+        tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path, subfolder="tokenizer", use_fast=False)
+
+    # support for Altdiffusion
+    text_encoder_cls = import_model_class_from_model_name_or_path(args.pretrained_model_name_or_path)
 
     # Load models and create wrapper for stable diffusion
     if args.train_only_unet:
       if os.path.exists(str(args.output_dir+"/text_encoder_trained")):
-        text_encoder = CLIPTextModel.from_pretrained(args.output_dir, subfolder="text_encoder_trained")
+        # text_encoder = CLIPTextModel.from_pretrained(args.output_dir, subfolder="text_encoder_trained")
+        text_encoder = text_encoder_cls.from_pretrained(args.output_dir, subfolder="text_encoder_trained")
       elif os.path.exists(str(args.output_dir+"/text_encoder")):
-        text_encoder = CLIPTextModel.from_pretrained(args.output_dir, subfolder="text_encoder")
+        # text_encoder = CLIPTextModel.from_pretrained(args.output_dir, subfolder="text_encoder")
+        text_encoder = text_encoder_cls.from_pretrained(args.output_dir, subfolder="text_encoder")
       else:
-        text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
+        # text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
+        text_encoder = text_encoder_cls.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
     else:
-      text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
+      # text_encoder = CLIPTextModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
+      text_encoder = text_encoder_cls.from_pretrained(args.pretrained_model_name_or_path, subfolder="text_encoder")
     vae = AutoencoderKL.from_pretrained(args.pretrained_model_name_or_path, subfolder="vae")
     unet = UNet2DConditionModel.from_pretrained(args.pretrained_model_name_or_path, subfolder="unet")
 
@@ -796,7 +820,7 @@ def run_training(args_imported):
                 if os.path.exists(frz_dir):
                   subprocess.call('rm -r '+ frz_dir, shell=True)
                 os.mkdir(frz_dir)
-                pipeline = StableDiffusionPipeline.from_pretrained(
+                pipeline = DiffusionPipeline.from_pretrained(
                     args.pretrained_model_name_or_path,
                     unet=accelerator.unwrap_model(unet),
                     text_encoder=accelerator.unwrap_model(text_encoder),
@@ -816,7 +840,7 @@ def run_training(args_imported):
                   print(" SAVING CHECKPOINT: "+args.Session_dir+"/"+inst+".ckpt")
                   # Create the pipeline using the trained modules and save it.
                   if accelerator.is_main_process:
-                     pipeline = StableDiffusionPipeline.from_pretrained(
+                     pipeline = DiffusionPipeline.from_pretrained(
                            args.pretrained_model_name_or_path,
                            unet=accelerator.unwrap_model(unet),
                            text_encoder=accelerator.unwrap_model(text_encoder),
@@ -839,7 +863,7 @@ def run_training(args_imported):
          txt_dir=args.output_dir + "/text_encoder_trained"
          if not os.path.exists(txt_dir):
            os.mkdir(txt_dir)
-         pipeline = StableDiffusionPipeline.from_pretrained(
+         pipeline = DiffusionPipeline.from_pretrained(
              args.pretrained_model_name_or_path,
              unet=accelerator.unwrap_model(unet),
              text_encoder=accelerator.unwrap_model(text_encoder),
@@ -847,7 +871,7 @@ def run_training(args_imported):
          pipeline.text_encoder.save_pretrained(txt_dir)       
 
       elif args.train_only_unet:
-        pipeline = StableDiffusionPipeline.from_pretrained(
+        pipeline = DiffusionPipeline.from_pretrained(
             args.pretrained_model_name_or_path,
             unet=accelerator.unwrap_model(unet),
             text_encoder=accelerator.unwrap_model(text_encoder),
@@ -857,7 +881,7 @@ def run_training(args_imported):
         subprocess.call('rm -r '+txt_dir, shell=True)
      
       else:
-        pipeline = StableDiffusionPipeline.from_pretrained(
+        pipeline = DiffusionPipeline.from_pretrained(
             args.pretrained_model_name_or_path,
             unet=accelerator.unwrap_model(unet),
             text_encoder=accelerator.unwrap_model(text_encoder),