22

app.py

@@ -6,6 +6,8 @@ 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 utils.image_utils import rescale_image
 from utils.prompt_utils import polish_prompt
@@ -13,11 +15,18 @@ from utils.prompt_utils import polish_prompt
 MODEL_REPO = "NewBie-AI/NewBie-image-Exp0.1"
 MAX_SEED = np.iinfo(np.int32).max
 
-pipe = NewbiePipeline.from_pretrained(
-    MODEL_REPO,
-    torch_dtype=torch.bfloat16,
-)
-pipe.to("cuda")
+transformer = NextDiT_3B_GQA_patch2_Adaln_Refiner_WHIT_CLIP.from_pretrained( 
+    MODEL_REPO , 
+    subfolder = 'transformer', 
+    torch_dtype=torch.bfloat16, 
+    use_safetensors=True
+).to("cuda")
+
+# pipe = NewbiePipeline.from_pretrained(
+#     MODEL_REPO,
+#     torch_dtype=torch.bfloat16,
+# )
+# pipe.to("cuda")
 
 # def prepare(prompt, is_polish_prompt):
 #     if not is_polish_prompt: return prompt, False

requirements.txt

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