chroma

app.py

@@ -3,7 +3,7 @@ import torch
 import gradio as gr
 import os
 from diffusers import FluxPipeline
-from transformers import T5EncoderModel
+from transformers import T5EncoderModel, CLIPTextModel
 
 # Initialize model outside the function
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -12,6 +12,13 @@ file_url = "https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
 
+# Initialize text encoders
+text_encoder = CLIPTextModel.from_pretrained(
+    single_file_base_model,
+    subfolder="text_encoder",
+    torch_dtype=dtype,
+    token=huggingface_token
+)
 
 text_encoder_2 = T5EncoderModel.from_pretrained(
     single_file_base_model,
@@ -24,11 +31,10 @@ text_encoder_2 = T5EncoderModel.from_pretrained(
 # Load the pipeline with proper configuration
 flux_pipeline = FluxPipeline.from_single_file(
     file_url,
+    text_encoder=text_encoder,
     text_encoder_2=text_encoder_2,
     torch_dtype=dtype,
-    token=huggingface_token,
-    use_safetensors=True,
-    variant="fp16"
+    token=huggingface_token
 )
 flux_pipeline.to(device)