Remove unnecessary import and minor style change

This commit is contained in:
Xrvk 2024-08-18 01:59:30 +03:00 committed by GitHub
parent 2e6ae1240c
commit 90af574794
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -15,7 +15,6 @@ from .layers import (
)
from einops import rearrange, repeat
import numpy as np
import comfy.ldm.common_dit
@dataclass
@ -33,6 +32,7 @@ class FluxParams:
qkv_bias: bool
guidance_embed: bool
class Flux(nn.Module):
"""
Transformer model for flow matching on sequences.
@ -106,9 +106,9 @@ class Flux(nn.Module):
if self.params.guidance_embed:
if guidance is None:
raise ValueError("Didn't get guidance strength for guidance distilled model.")
vec.add_(self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype)))
vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
vec.add_(self.vector_in(y))
vec = vec + self.vector_in(y)
txt = self.txt_in(txt)
ids = torch.cat((txt_ids, img_ids), dim=1)