How to use the UNet2DConditionModel to pick a class?

[rsousa95]

Hi all,

I'm a bit new to these things and I'm trying to train some diffusion models. I have 10 classes and I'd like to be able to ask the class I want. From what I've seen in most tutorials, this model is used with text prompts, but I don't need something as complicated, I think.

This tutorial from the Hugging Face course instead uses the regular unconditional UNet2DModel and uses the class as an extra channel in the input noise. If the UNet2DConditionModel could accept a class, I'd rather use that.

[tolgacangoz]

To understand better I need to ask: Could you please explain why you need UNet2DConditionModel? UNet2DConditionModel requires to see text, otherwise it cannot work.

[asomoza]

I’m not that well-versed in the Pokemon world, but for example, for the fire attribute, you’ll need to caption all that doesn’t imply fire in the picture and vice versa.

For instance, if you have a picture of a charmander in the forest with grass and sky as the background you should include all the background details and the details of the Charmander that you don’t consider unique to the fire attribute. This could include the color of its eyes, whether it resembles a dinosaur or reptile, its skin color, etc. However, if you want all the fire attribute Pokemon to have a flame, you shouldn’t caption that because it would be implied in the fire keyword.

Also you'll need a big dataset to make the model kind of forget what a fire is, because if not, all the pokemons will appear under a fire, burning, etc.

[Conscht]

For that, you're going to need the complicated thing that you said, if you want to be able to ask the model for something you need to train with text, and actually when you say "fire" of "water" you're using a text prompt.

It's easy and fast to train a dreambooth LoRA for example with just one specific pokemon, but I for what you're asking you'll need to do a full finetune. or maybe train a bunch of dreambooth LoRAs an try to merge them to one model.

So actually for what you're asking is more work than what most people do, you can find the guide here: https://huggingface.co/docs/diffusers/training/text2image where they actually train with the pokemon dataset (that doesn't exist anymore).

You'll need a decent sized dataset and if you caption all the pokemons with fire, water, earth, etc, you can ask the model to give you an image of that type of pokemon, how good the model learns those concepts depends on the size of the dataset and how good are the captions for them, for example just writing fire on the caption won't work.

Hello, so the Unet2DConditional only works with Text embeddings? I thought you might could replace the Text embeddings /encoder hidden state with custome ones?

[asomoza]

Hi, I didn't say you couldn't use something else, I was just answering the question where the OP wanted to ask the model for a class as a text e.g. "fire", "water", etc.

[Conscht]

It was more like a question, because im unsure about it, and try to implement it in my architecture.

I have the Unet, concatenated two embeddings (Identity Embedding and Emotion Embedding), and pass it to the encoder hidden states. Is that a valid approach?

When i only conditioned it on one emebdding, the results where not good, nor terrible bad. After I added another embedding, the results became way worse. What could I do?

[asomoza]

I'm not an expert at training and sadly I still haven't had the time to train a UNet2DConditionModel but in theory for what I know it should work. Maybe its just a matter of training data or training time.

Same as OP, it probably should have worked with just the Unet2DModel and the class as an extra channel but since the dataset was a lot more complex than some black and white numbers and letters, and also with some "hidden" class attribute, the model probably needs a lot more training data/epochs than just using a base UNet2DConditionModel and finetune it with text prompts.

Maybe you can compare your results with the first releases of Stable Diffusion, we got to where we are now because they scaled up the training data and the model parameters.

But overall it makes kind of sense that the results got worse with the second embedding if you didn't add a lot more training to the model.

[a-caycioglu]

What are your class_embed_type and num_class_embeds hyperparameters? I think these two should be adjusted in your case.
model = UNet2DConditionModel( sample_size=64, in_channels=3, out_channels=3, layers_per_block=2, block_out_channels=(128, 256, 512, 512), down_block_types=("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D"), up_block_types=("UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D"), class_embed_type="projection", # You can experiment with different types of arguments num_class_embeds=10 # Your 10 classes )