chroma

app.py

@@ -2,22 +2,51 @@ import spaces
 import torch
 import gradio as gr
 import os
-from diffusers import FluxPipeline
+from diffusers import FluxPipeline, FluxTransformer2DModel, BitsAndBytesConfig
+from transformers import T5EncoderModel, BitsAndBytesConfig as BitsAndBytesConfigTF
 
 # Initialize model outside the function
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16
 file_url = "https://huggingface.co/lodestones/Chroma/blob/main/chroma-unlocked-v31.safetensors"
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
+single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
 
-# Load the pipeline with proper configuration
-flux_pipeline = FluxPipeline.from_single_file(
+# Initialize text encoder
+quantization_config_tf = BitsAndBytesConfigTF(load_in_8bit=True, bnb_8bit_compute_dtype=torch.bfloat16)
+text_encoder_2 = T5EncoderModel.from_pretrained(
+    single_file_base_model,
+    subfolder="text_encoder_2",
+    torch_dtype=dtype,
+    config=single_file_base_model,
+    quantization_config=quantization_config_tf,
+    token=huggingface_token
+)
+
+# Initialize transformer
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    token=huggingface_token
+)
+transformer = FluxTransformer2DModel.from_single_file(
     file_url,
+    subfolder="transformer",
+    torch_dtype=dtype,
+    config=single_file_base_model,
+    quantization_config=quantization_config,
+    token=huggingface_token
+)
+
+# Load the pipeline with proper configuration
+flux_pipeline = FluxPipeline.from_pretrained(
+    single_file_base_model,
+    transformer=transformer,
+    text_encoder_2=text_encoder_2,
     torch_dtype=dtype,
-    token=huggingface_token,
-    use_safetensors=True,
-    local_files_only=False,
-    config_file="model_index.json"
+    token=huggingface_token
 )
 flux_pipeline.to(device)