22

app.py

@@ -6,8 +6,7 @@ import torch
 # from transformers import AutoProcessor, AutoTokenizer
 # from diffusers import DiffusionPipeline
 from diffusers import NewbiePipeline
-from diffusers.models.transformers.transformer_newbie import NextDiT_3B_GQA_patch2_Adaln_Refiner_WHIT_CLIP
-from diffusers.models.transformers.transformer_newbie import NextDiT_CLIP as nextdit_clip
+from transformers import AutoModel
 
 
 from utils.image_utils import rescale_image
@@ -16,12 +15,11 @@ from utils.prompt_utils import polish_prompt
 MODEL_REPO = "NewBie-AI/NewBie-image-Exp0.1"
 MAX_SEED = np.iinfo(np.int32).max
 
-transformer = NextDiT_3B_GQA_patch2_Adaln_Refiner_WHIT_CLIP.from_pretrained( 
-    MODEL_REPO,
-    subfolder = 'transformer',
-    torch_dtype=torch.bfloat16, 
-    use_safetensors=True
-).to("cuda")
+device = "cuda"
+model_path = "Disty0/NewBie-image-Exp0.1-Diffusers"
+text_encoder_2 = AutoModel.from_pretrained(model_path, subfolder="text_encoder_2", trust_remote_code=True, torch_dtype=torch.bfloat16)
+pipe = NewbiePipeline.from_pretrained(model_path, text_encoder_2=text_encoder_2, torch_dtype=torch.bfloat16)
+del text_encoder_2
 
 # pipe = NewbiePipeline.from_pretrained(
 #     MODEL_REPO,

requirements.txt

@@ -4,4 +4,4 @@ transformers
 accelerate
 spaces
 flash-attn
-git+https://github.com/E-Anlia/diffusers.git@add-newbie-pipeline
+git+https://github.com/Disty0/diffusers