chroma

app.py

@@ -2,13 +2,13 @@ import spaces
 import torch
 import gradio as gr
 import os
-from diffusers import FluxPipeline, FluxTransformer2DModel, BitsAndBytesConfig
+from diffusers import FluxPipeline
 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"
+file_url = "https://huggingface.co/lodestones/Chroma/resolve/main/chroma-unlocked-v31.safetensors"
 huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
 single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
 
@@ -23,30 +23,14 @@ text_encoder_2 = T5EncoderModel.from_pretrained(
     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,
+flux_pipeline = FluxPipeline.from_single_file(
+    file_url,
     text_encoder_2=text_encoder_2,
     torch_dtype=dtype,
-    token=huggingface_token
+    token=huggingface_token,
+    use_safetensors=True,
+    variant="fp16"
 )
 flux_pipeline.to(device)