twin model

app.py

@@ -1,11 +1,10 @@
 import gradio as gr
 import numpy as np
 import random
-import spaces
 import torch
 from PIL import Image
 from torchvision import transforms
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, AutoencoderKL
 
 # Define constants
 dtype = torch.bfloat16
@@ -13,7 +12,11 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Load the diffusion pipeline with optimizations
+# Load the initial VAE model for preprocessing
+vae_model_name = "CompVis/stable-diffusion-v1-4"  # Example VAE model
+vae = AutoencoderKL.from_pretrained(vae_model_name).to(device)
+
+# Load the FLUX diffusion pipeline with optimizations
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype)
 pipe.enable_model_cpu_offload()
 pipe.vae.enable_slicing()
@@ -27,11 +30,13 @@ def preprocess_image(image, image_size):
         transforms.Normalize([0.5], [0.5])
     ])
     image = preprocess(image).unsqueeze(0).to(device, dtype=dtype)
+    print("Image processed successfully.")
     return image
 
 def encode_image(image, vae):
     with torch.no_grad():
         latents = vae.encode(image).latent_dist.sample() * 0.18215
+    print("Image encoded successfully.")
     return latents
 
 @spaces.GPU()
@@ -61,20 +66,33 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             return fallback_image, seed
     else:
         # img2img case
-        vae_image_size = pipe.vae.config.sample_size  # Ensure this is correct
+        print("Initial image provided, starting preprocessing...")
+        vae_image_size = 1024  # Using FLUX VAE sample size for preprocessing
         init_image = init_image.convert("RGB")
         init_image = preprocess_image(init_image, vae_image_size)
-        latents = encode_image(init_image, pipe.vae)
-
+        
+        print("Starting encoding of the image...")
+        latents = encode_image(init_image, vae)
+
+        print(f"Latents shape after encoding: {latents.shape}")
+        
+        # Ensure the latents size matches the expected input size for the FLUX model
+        print("Interpolating latents to match model's input size...")
         latents = torch.nn.functional.interpolate(latents, size=(height // 8, width // 8))
-        latent_channels = pipe.vae.config.latent_channels  # Ensure this is correct
-        if latent_channels != 64:
-            conv = torch.nn.Conv2d(latent_channels, 64, kernel_size=1).to(device, dtype=dtype)
+        
+        latent_channels = 16  # Using FLUX VAE latent channels
+        print(f"Latent channels from VAE: {latent_channels}, expected by FLUX model: {pipe.vae.config.latent_channels}")
+        
+        if latent_channels != pipe.vae.config.latent_channels:
+            print(f"Adjusting latent channels from {latent_channels} to {pipe.vae.config.latent_channels}")
+            conv = torch.nn.Conv2d(latent_channels, pipe.vae.config.latent_channels, kernel_size=1).to(device, dtype=dtype)
             latents = conv(latents)
 
-        latents = latents.permute(0, 2, 3, 1).contiguous().view(-1, 64)
+        latents = latents.permute(0, 2, 3, 1).contiguous().view(-1, pipe.vae.config.latent_channels)
+        print(f"Latents shape after permutation: {latents.shape}")
 
         try:
+            print("Sending latents to the FLUX transformer...")
             # Determine if 'timesteps' is required for the transformer
             if hasattr(pipe.transformer, 'forward') and hasattr(pipe.transformer.forward, '__code__') and 'timesteps' in pipe.transformer.forward.__code__.co_varnames:
                 timestep = torch.tensor([num_inference_steps], device=device, dtype=dtype)
@@ -86,6 +104,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             return fallback_image, seed
 
         try:
+            print("Generating final image with the FLUX pipeline...")
             image = pipe(
                 prompt=prompt,
                 height=height,
@@ -95,6 +114,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
                 guidance_scale=0.0,
                 latents=latents
             ).images[0]
+            print("Image generation completed.")
         except Exception as e:
             print(f"Pipeline call with latents failed with error: {e}")
             return fallback_image, seed
@@ -210,9 +230,8 @@ with gr.Blocks(css=css) as demo:
             cache_examples="lazy"
         )
 
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
+    run_button.click(
+        infer,
         inputs=[prompt, init_image, seed, randomize_seed, width, height, num_inference_steps],
         outputs=[result, seed]
     )