added console logging

app.py

@@ -17,6 +17,7 @@ MAX_IMAGE_SIZE = 2048
 pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
 
 def preprocess_image(image, image_size):
+    print(f"Preprocessing image to size: {image_size}x{image_size}")
     # Preprocess the image for the VAE
     preprocess = transforms.Compose([
         transforms.Resize((image_size, image_size)),  # Use model-specific size
@@ -24,47 +25,58 @@ def preprocess_image(image, image_size):
         transforms.Normalize([0.5], [0.5])  # Ensure this matches the VAE's training normalization
     ])
     image = preprocess(image).unsqueeze(0).to(device, dtype=dtype)
+    print(f"Image shape after preprocessing: {image.shape}")
     return image
 
 def encode_image(image, vae):
+    print("Encoding image using the VAE")
     # Encode the image using the VAE
     with torch.no_grad():
         latents = vae.encode(image).latent_dist.sample() * 0.18215
+    print(f"Latents shape after encoding: {latents.shape}")
     return latents
 
 @spaces.GPU()
 def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=4, progress=gr.Progress(track_tqdm=True)):
+    print(f"Inference started with prompt: {prompt}")
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
+    print(f"Using seed: {seed}")
     generator = torch.Generator().manual_seed(seed)
 
     # Get the expected image size for the VAE
     vae_image_size = pipe.vae.config.sample_size
+    print(f"Expected VAE image size: {vae_image_size}")
     
     if init_image is not None:
-        # Process img2img
+        print("Initial image provided, processing img2img")
         init_image = init_image.convert("RGB")
         init_image = preprocess_image(init_image, vae_image_size)
         latents = encode_image(init_image, pipe.vae)
 
-        # Debug: Print the shape of the latents after encoding
-        print(f"Latents shape after encoding: {latents.shape}")
-
-        # Ensure latents are correctly shaped and adjusted
+        # Interpolating latents
+        print(f"Interpolating latents to size: {(height // 8, width // 8)}")
         latents = torch.nn.functional.interpolate(latents, size=(height // 8, width // 8))
+        print(f"Latents shape after interpolation: {latents.shape}")
         
         # Convert latent channels to 64 as expected by the transformer
         latent_channels = pipe.vae.config.latent_channels
+        print(f"Expected latent channels: 64, current latent channels: {latent_channels}")
         if latent_channels != 64:
+            print(f"Converting latent channels from {latent_channels} to 64")
             conv = torch.nn.Conv2d(latent_channels, 64, kernel_size=1).to(device, dtype=dtype)
             latents = conv(latents)
+            print(f"Latents shape after channel conversion: {latents.shape}")
 
         # Reshape latents to match the transformer's input expectations
         latents = latents.view(1, 64, height // 8, width // 8)
-
-        # Debug: Print the shape of the latents after reshaping
         print(f"Latents shape after reshaping: {latents.shape}")
 
+        # Flatten the latents if required by the transformer
+        latents = latents.flatten(start_dim=1)
+        print(f"Latents shape after flattening: {latents.shape}")
+
+        print("Calling the diffusion pipeline with latents")
         image = pipe(
             prompt=prompt,
             height=height,
@@ -75,7 +87,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             latents=latents
         ).images[0]
     else:
-        # Process text2img
+        print("No initial image provided, processing text2img")
         image = pipe(
             prompt=prompt,
             height=height,
@@ -85,10 +97,12 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             guidance_scale=0.0
         ).images[0]
     
+    print("Inference complete")
     return image, seed
 
 
 
+
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",