Files changed (1) hide show
  1. app.py +7 -5
app.py CHANGED
@@ -274,18 +274,20 @@ with gr.Blocks() as demo:
274
  # Initialize both pipelines
275
  if __name__ == "__main__":
276
  from diffusers import FluxTransformer2DModel, FluxPipeline, BitsAndBytesConfig
 
277
  # Initialize Flux pipeline
278
  device = "cuda" if torch.cuda.is_available() else "cpu"
279
  huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
280
 
281
- quantization_config = BitsAndBytesConfig(load_in_4bit=True)
282
- hf_token = ""
283
  dtype = torch.bfloat16
284
  file_url = "https://huggingface.co/gokaygokay/flux-game/blob/main/gokaygokay_00001_.safetensors"
285
  single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
286
- transformer = FluxTransformer2DModel.from_single_file(file_url, subfolder="transformer", torch_dtype=dtype, config=single_file_base_model, quantization_config=quantization_config)
287
- flux_pipeline = FluxPipeline.from_pretrained(single_file_base_model, transformer=transformer, torch_dtype=dtype, token=huggingface_token, quantization_config=quantization_config)
288
-
 
 
 
289
  # Initialize Trellis pipeline
290
  trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
291
  trellis_pipeline.cuda()
 
274
  # Initialize both pipelines
275
  if __name__ == "__main__":
276
  from diffusers import FluxTransformer2DModel, FluxPipeline, BitsAndBytesConfig
277
+ from transformers import BitsAndBytesConfig as BitsAndBytesConfigTF
278
  # Initialize Flux pipeline
279
  device = "cuda" if torch.cuda.is_available() else "cpu"
280
  huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
281
 
 
 
282
  dtype = torch.bfloat16
283
  file_url = "https://huggingface.co/gokaygokay/flux-game/blob/main/gokaygokay_00001_.safetensors"
284
  single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
285
+ quantization_config_tf = BitsAndBytesConfigTF(load_in_4bit=True)
286
+ 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)
287
+ quantization_config = BitsAndBytesConfig(load_in_4bit=True)
288
+ transformer = FluxTransformer2DModel.from_single_file(file_url, subfolder="transformer", torch_dtype=dtype, config=single_file_base_model, quantization_config=quantization_config, token=huggingface_token)
289
+ flux_pipeline = FluxPipeline.from_pretrained(single_file_base_model, transformer=transformer, text_encoder_2=text_encoder_2, torch_dtype=dtype, quantization_config=quantization_config, token=huggingface_token)
290
+
291
  # Initialize Trellis pipeline
292
  trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")
293
  trellis_pipeline.cuda()