mapped projection

app.py

@@ -13,6 +13,9 @@ dtype = torch.bfloat16
 device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
+LATENT_CHANNELS = 16
+TRANSFORMER_IN_CHANNELS = 64
+SCALING_FACTOR = 0.3611
 
 # Load FLUX model
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
@@ -20,8 +23,8 @@ pipe.enable_model_cpu_offload()
 pipe.vae.enable_slicing()
 pipe.vae.enable_tiling()
 
-# Add a projection layer to match x_embedder input
-projection = nn.Linear(16, 64).to(device).to(dtype)
+# Add a projection layer to match transformer input
+projection = nn.Linear(LATENT_CHANNELS, TRANSFORMER_IN_CHANNELS).to(device).to(dtype)
 
 def preprocess_image(image, image_size):
     preprocess = transforms.Compose([
@@ -41,7 +44,7 @@ def process_latents(latents, height, width):
     print(f"Latent shape after potential interpolation: {latents.shape}")
     
     # Reshape latents to [batch_size, seq_len, channels]
-    latents = latents.permute(0, 2, 3, 1).reshape(1, -1, latents.shape[1])
+    latents = latents.permute(0, 2, 3, 1).reshape(1, -1, LATENT_CHANNELS)
     print(f"Reshaped latent shape: {latents.shape}")
     
     # Project latents from 16 to 64 dimensions
@@ -73,10 +76,10 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             init_image = preprocess_image(init_image, 1024)  # Using 1024 as FLUX VAE sample size
             
             # Encode the image using FLUX VAE
-            latents = pipe.vae.encode(init_image).latent_dist.sample() * 0.18215
+            latents = pipe.vae.encode(init_image).latent_dist.sample() * SCALING_FACTOR
             print(f"Initial latent shape from VAE: {latents.shape}")
             
-            # Process latents to match x_embedder input
+            # Process latents to match transformer input
             latents = process_latents(latents, height, width)
 
             print(f"x_embedder weight shape: {pipe.transformer.x_embedder.weight.shape}")