22

app.py

@@ -17,8 +17,18 @@ MAX_SEED = np.iinfo(np.int32).max
 
 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)
+text_encoder_2 = AutoModel.from_pretrained(
+    model_path, 
+    subfolder="text_encoder_2", 
+    trust_remote_code=True, 
+    torch_dtype=torch.bfloat16,
+    device_map="cuda",
+)
+pipe = NewbiePipeline.from_pretrained(
+    model_path, 
+    text_encoder_2=text_encoder_2, 
+    torch_dtype=torch.bfloat16
+).to("cuda")
 del text_encoder_2
 
 

requirements.txt

@@ -3,7 +3,7 @@ torch
 transformers
 accelerate
 spaces
-flash-attn
+# flash-attn
 git+https://github.com/Disty0/diffusers
 timm
 torchvision