How to use flux.1-dev models in python?

[asatulsinha4]

Hi All

I am trying to use flux dev models locally, but unable to. There are 2 files, one that is only unet and second with is all in one with text encoders and vae included.
So far, I've been running flux dev and schnell locally using huggingface diffusers lib, and recently started exploring LoRAs and other flux models available on the internet.

This is how I've been using so far:

import torch
from diffusers import FluxPipeline
from datetime import datetime

prompt = "an ancient medeival europe city market, with a large crowd of people, realism, sun, soft color, balanced colors, 4K marterpiece, detailed"

now = datetime.now()
hour, minute, second = now.hour % 24 ,now.minute,now.second
local_timestamp = f"{now.date()}_{hour}_{minute:02d}_{second:02}" 

pipe = FluxPipeline.from_pretrained("E:/Projects/flux_lora/FLUX.1-dev", torch_dtype=torch.bfloat16)
full_image_path =  'E:\\Projects\\flux_lora\\images\\d_i_{0}'.format(local_timestamp)

pipe.enable_sequential_cpu_offload() #this does a better job for me

images = pipe(
    prompt = prompt,
    height=512,
    width=512,
    guidance_scale=3.5,
    num_inference_steps=28, # for dev: 28 (50 is better), for schnell: 4
    # max_sequence_length=256,  # for dev: 512, for schnell: 256
    num_images_per_prompt=5,
    # generator=torch.Generator("cpu").manual_seed(0)
).images


for idx, image in enumerate(images):
    image_path = f"{full_image_path}_{idx + 1}.png"
    # image_path = os.path.join(output_dir, f"image_{local_timestamp}_{idx + 1}.png")
    image.save(image_path, format="PNG")
    print(f"Saved image to {image_path}")
    image.show()

I have resources constraint. I run it on my laptop, which has core i9 11900H CPU, 32GB RAM, and NVIDIA RTX 3060 Laptop GPU with 6GB of VRAM. Also one of the reasons why I am unwilling to explore popular options like ComfyUI, Automatic 1111 etc

Although this is slow, but it generates images within 4 minutes, so I'm fine with it.

I recently came across an flux model, which I'm really dying to try, but for the life of me I am unable to figure out how to make it work.
I've tried searching it online, using chatgpt...no help.

Also, I am a complete novice in python. So please feel free to take me as an idiot.

I come from Java and Perl background.

I've tried this:
pipeline = StableDiffusionPipeline.from_single_file(aio_safetensors_path, use_safetensors=True)
and this:

pipe = FluxPipeline.from_pretrained("E:/Projects/flux_lora/FLUX.1-dev", torch_dtype=torch.bfloat16)
unet_config = UNet2DConditionModel.load_config(unet_path)

custom_unet = UNet2DConditionModel(**unet_config)

state_dict = load_file(unet_path)

custom_unet.load_state_dict(state_dict)

pipe.unet = custom_unet

and this:

transformer = FluxTransformer2DModel.from_single_file(unet_safetensors_path)
 # for dev
pipeline = FluxPipeline.from_pretrained("E:/Projects/flux_lora/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16)

In every case, it errors out like this:

Traceback (most recent call last):
  File "e:\Projects\flux_lora\u.ai\u_custom_aio.py", line 23, in <module>
    transformer = FluxTransformer2DModel.from_single_file(unet_safetensors_path)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "E:\Projects\flux_lora\.venv\Lib\site-packages\huggingface_hub\utils\_validators.py", line 114, in _inner_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "E:\Projects\flux_lora\.venv\Lib\site-packages\diffusers\loaders\single_file_model.py", line 291, in from_single_file
    config = fetch_diffusers_config(checkpoint)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "E:\Projects\flux_lora\.venv\Lib\site-packages\diffusers\loaders\single_file_utils.py", line 643, in fetch_diffusers_config
    model_type = infer_diffusers_model_type(checkpoint)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "E:\Projects\flux_lora\.venv\Lib\site-packages\diffusers\loaders\single_file_utils.py", line 598, in infer_diffusers_model_type
    if checkpoint["img_in.weight"].shape[1] == 384:
       ~~~~~~~~~~^^^^^^^^^^^^^^^^^
KeyError: 'img_in.weight'

I only have safetensors files and no config json files, and it should be able to run on automatic 1111 like this

can someone help please? I want to make it work using python, preferably by using huggingface diffusers and safetensors libraries

Thanks!!