Update parallax.py

parallax.py

@@ -34,7 +34,7 @@ LAMA_TRAIN_CFG = OmegaConf.load('big-lama/config.yaml')
 LAMA_TRAIN_CFG['training_model']['predict_only'] = True
 LAMA = load_checkpoint(LAMA_TRAIN_CFG, 'big-lama/models/best.ckpt', strict=False, map_location='cpu')
 LAMA = LAMA.to('cuda' if torch.cuda.is_available() else 'cpu').eval()
-FLUX_KONTEXT_INPAINT_PIPELINE = FluxKontextInpaintPipeline.from_pretrained('black-forest-labs/FLUX.1-Kontext-dev', dtype=torch.bfloat16, token=os.environ['HF_TOKEN']).to('cuda' if torch.cuda.is_available() else 'cpu')
+FLUX_KONTEXT_INPAINT_PIPELINE = FluxKontextInpaintPipeline.from_pretrained('black-forest-labs/FLUX.1-Kontext-dev', torch_dtype=torch.bfloat16, token=os.environ['HF_TOKEN']).to('cuda' if torch.cuda.is_available() else 'cpu')
 FLUX_KONTEXT_INPAINT_PIPELINE.load_lora_weights('alimama-creative/FLUX.1-Turbo-Alpha')
 FLUX_KONTEXT_INPAINT_PIPELINE.fuse_lora()
 optimize_pipeline_(FLUX_KONTEXT_INPAINT_PIPELINE, image=Image.new('RGB', (1024, 1024)), mask_image=Image.new('L', (512, 512)), prompt='prompt')