monster/Qwen-Image-Layered-MRP-MLX
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Do text encoding directly and remove TextEncoder class

Browse Source
This commit is contained in:
filipstrand 2024-08-15 20:55:56 +02:00
parent f35eec2a3b
commit 0fae839053
2 changed files with 9 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
src/flux_1_schnell/models/flux.py
View File

@ -3,18 +3,16 @@ import mlx.core as mx
from PIL import Image
from tqdm import tqdm
from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
from flux_1_schnell.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
from flux_1_schnell.config.config import Config
from flux_1_schnell.latent_creator.latent_creator import LatentCreator
from flux_1_schnell.models.text_encoder.text_encoder import TextEncoder
from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from flux_1_schnell.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from flux_1_schnell.models.transformer.transformer import Transformer
from flux_1_schnell.models.vae.vae import VAE
from flux_1_schnell.post_processing.image_util import ImageUtil
from flux_1_schnell.scheduler.scheduler import FlowMatchEulerDiscreteNoiseScheduler
from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
from flux_1_schnell.tokenizer.tokenizer_handler import TokenizerHandler
from flux_1_schnell.weights.weight_handler import WeightHandler
@ -34,13 +32,11 @@ class Flux1Schnell:
def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image:
latents = LatentCreator.create(seed)
prompt_embeds, pooled_prompt_embeds = TextEncoder.encode(
prompt=prompt,
clip_tokenizer=self.clip_tokenizer,
t5_tokenizer=self.t5_tokenizer,
clip_text_encoder=self.clip_text_encoder,
t5_text_encoder=self.t5_text_encoder
)
t5_tokens = self.t5_tokenizer.tokenize(prompt)
clip_tokens = self.clip_tokenizer.tokenize(prompt)
prompt_embeds = self.t5_text_encoder.forward(t5_tokens)
pooled_prompt_embeds = self.clip_text_encoder.forward(clip_tokens)
for t in tqdm(range(config.num_inference_steps)):
noise = self.transformer.predict(

25
src/flux_1_schnell/models/text_encoder/text_encoder.py
View File

@ -1,25 +0,0 @@
import mlx.core as mx
from flux_1_schnell.models.text_encoder.clip_encoder.clip_encoder import CLIPEncoder
from flux_1_schnell.models.text_encoder.t5_encoder.t5_encoder import T5Encoder
from flux_1_schnell.tokenizer.clip_tokenizer import TokenizerCLIP
from flux_1_schnell.tokenizer.t5_tokenizer import TokenizerT5
class TextEncoder:
@staticmethod
def encode(
prompt: str,
clip_tokenizer: TokenizerCLIP,
t5_tokenizer: TokenizerT5,
clip_text_encoder: CLIPEncoder,
t5_text_encoder: T5Encoder
) -> (mx.array, mx.array):
clip_tokens = clip_tokenizer.tokenize(prompt)
pooled_prompt_embeds = clip_text_encoder.forward(clip_tokens)
t5_tokens = t5_tokenizer.tokenize(prompt)
prompt_embeds = t5_text_encoder.forward(t5_tokens)
return prompt_embeds, pooled_prompt_embeds