twodgirl/flux-qwen-neutered

Flux Qwen Neutered

Qwen is a lightweight alternative to the T5 model. For use with the Flux Dev model.

This repo is an experimental work, and not a final replacement for the built-in text encoder.

Inference

from diffusers import FluxPipeline, FluxTransformer2DModel
from text_encoder import PretrainedTextEncoder
import torch
from typing import List, Optional, Union

class FluxQwenPipeline(FluxPipeline):
    def _get_t5_prompt_embeds(self,
                              prompt: Union[str, List[str]] = None,
                              num_images_per_prompt: int = 1,
                              max_sequence_length: int = 512,
                              device: Optional[torch.device] = None,
                              dtype: Optional[torch.dtype] = None):
        qwen_out, qwen_mask = self.encode_qwen(prompt, max_sequence_length, device)
        output = self.text_encoder_2(qwen_out, max_length=max_sequence_length)

        return output * qwen_mask

    def encode_qwen(self, prompt, max_sequence_length=256, device=None):
        inputs = self.qwen_tokenizer(prompt,
                                     return_tensors='pt',
                                     padding='max_length',
                                     truncation=True,
                                     max_length=max_sequence_length)
        input_ids = inputs['input_ids'].to(device)
        attention_mask = inputs['attention_mask'].to(device)
        output = self.qwen_model(input_ids=input_ids,
                                 attention_mask=attention_mask)

        return output.last_hidden_state, attention_mask.unsqueeze(-1)

if __name__ == '__main__':
    pipe = FluxQwenPipeline.from_pretrained('black-forest-labs/FLUX.1-dev',
                                            text_encoder_2=None,
                                            torch_dtype=torch.bfloat16)
    PretrainedTextEncoder.setup_pipeline(pipe,
                                         text_encoder_path='twodgirl/flux-qwen-neutered',
                                         device='cuda')
    pipe.enable_model_cpu_offload()
    image = pipe('a black cat wearing a Pikachu cosplay').images[0]
    image.save('cat.png')

Datasets

• alfredplpl/artbench-pd-256x256
• BaiqiL/GenAI-Bench
• data-is-better-together/open-image-preferences-v1-binarized
• kadirnar/fluxdev_controlnet_16k
• NewEden/CivitAI-SD-Prompts
• sengunsipahi/civit_0.1_p1
• SPRIGHT-T2I/spright_coco
• twodgirl/flux-lora-tag-explorer
• twodgirl/spatial-caption

Disclaimer

Use of this code and the model requires citation and attribution to the author via a link to their Hugging Face profile in all resulting work.