chroma

app.py

@@ -7,16 +7,15 @@ from transformers import T5EncoderModel, BitsAndBytesConfig as BitsAndBytesConfi
 # Initialize model outside the function
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16
-base_model = "lodestones/Chroma"
 file_url = "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v31.safetensors"
 
 quantization_config_tf = BitsAndBytesConfigTF(load_in_8bit=True, bnb_8bit_compute_dtype=torch.bfloat16)
-text_encoder_2 = T5EncoderModel.from_pretrained(base_model, subfolder="text_encoder_2", torch_dtype=dtype, config=base_model, quantization_config=quantization_config_tf)
+text_encoder_2 = T5EncoderModel.from_single_file(file_url, subfolder="text_encoder_2", torch_dtype=dtype, quantization_config=quantization_config_tf)
 
 quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16)
-transformer = FluxTransformer2DModel.from_single_file(file_url, subfolder="transformer", torch_dtype=dtype, config=base_model, quantization_config=quantization_config)
+transformer = FluxTransformer2DModel.from_single_file(file_url, subfolder="transformer", torch_dtype=dtype, quantization_config=quantization_config)
 
-flux_pipeline = FluxPipeline.from_pretrained(base_model, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype)
+flux_pipeline = FluxPipeline.from_single_file(file_url, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype)
 flux_pipeline.to(device)
 
 @spaces.GPU()