chroma

app.py

@@ -1,21 +1,14 @@
 import spaces
 import torch
 import gradio as gr
-from diffusers import FluxTransformer2DModel, FluxPipeline, BitsAndBytesConfig
-from transformers import T5EncoderModel, BitsAndBytesConfig as BitsAndBytesConfigTF
+from diffusers import FluxPipeline
 
 # Initialize model outside the function
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16
 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_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, quantization_config=quantization_config)
-
-flux_pipeline = FluxPipeline.from_single_file(file_url, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype)
+flux_pipeline = FluxPipeline.from_single_file(file_url, torch_dtype=dtype)
 flux_pipeline.to(device)
 
 @spaces.GPU()