spacesgpu

app.py

@@ -273,16 +273,18 @@ with gr.Blocks() as demo:
 
 # Initialize both pipelines
 if __name__ == "__main__":
+    from diffusers import FluxTransformer2DModel, FluxPipeline, BitsAndBytesConfig
     # Initialize Flux pipeline
     device = "cuda" if torch.cuda.is_available() else "cpu"
     huggingface_token = os.getenv("HUGGINGFACE_TOKEN")
-    flux_pipeline = DiffusionPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-dev", 
-        torch_dtype=torch.bfloat16,
-        token = huggingface_token
-    )
-    flux_pipeline.load_lora_weights("gokaygokay/Flux-Game-Assets-LoRA-v2")
-    flux_pipeline = flux_pipeline.to(device)
+    
+    quantization_config = BitsAndBytesConfig(load_in_4bit=True)
+    hf_token = ""
+    dtype = torch.bfloat16
+    file_url = "https://huggingface.co/gokaygokay/flux-game/blob/main/gokaygokay_00001_.safetensors"
+    single_file_base_model = "camenduru/FLUX.1-dev-diffusers"
+    transformer = FluxTransformer2DModel.from_single_file(file_url, subfolder="transformer", torch_dtype=dtype, config=single_file_base_model, quantization_config=quantization_config)
+    pipe = FluxPipeline.from_pretrained(single_file_base_model, transformer=transformer, torch_dtype=dtype, token=huggingface_token, quantization_config=quantization_config)
 
     # Initialize Trellis pipeline
     trellis_pipeline = TrellisImageTo3DPipeline.from_pretrained("JeffreyXiang/TRELLIS-image-large")