model compatibility

app.py

@@ -2,57 +2,65 @@ import gradio as gr
 import numpy as np
 import random
 import torch
-import spaces
 from PIL import Image
 from torchvision import transforms
 from diffusers import DiffusionPipeline, AutoencoderKL
+import spaces
 
 # Define constants
-dtype = torch.bfloat16
+flux_dtype = torch.bfloat16
+vae_dtype = torch.float32
 device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Load the initial VAE model for preprocessing YAY
-# Load the initial VAE model for preprocessing
-vae_model_name = "runwayml/stable-diffusion-v1-5"  # Adjusted VAE model path
-vae = AutoencoderKL.from_pretrained(vae_model_name, subfolder="vae").to(device)
-
+def load_models():
+    # Load the initial VAE model for preprocessing in float32
+    vae_model_name = "runwayml/stable-diffusion-v1-5"
+    vae = AutoencoderKL.from_pretrained(vae_model_name, subfolder="vae").to(device).to(vae_dtype)
+    
+    # Load the FLUX diffusion pipeline with bfloat16
+    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=flux_dtype)
+    pipe.enable_model_cpu_offload()
+    pipe.vae.enable_slicing()
+    pipe.vae.enable_tiling()
+    pipe.to(device)
+    
+    return vae, pipe
 
-# 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()
-pipe.vae.enable_tiling()
-pipe.to(device)
+vae, pipe = load_models()
 
 def preprocess_image(image, image_size):
     preprocess = transforms.Compose([
-        transforms.Resize((image_size, image_size)),
+        transforms.Resize((image_size, image_size), interpolation=transforms.InterpolationMode.LANCZOS),
         transforms.ToTensor(),
         transforms.Normalize([0.5], [0.5])
     ])
-    image = preprocess(image).unsqueeze(0).to(device, dtype=dtype)
+    image = preprocess(image).unsqueeze(0).to(device, dtype=vae_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
+    try:
+        with torch.no_grad():
+            latents = vae.encode(image).latent_dist.sample() * 0.18215
+        print("Image encoded successfully.")
+        return latents
+    except RuntimeError as e:
+        print(f"Error during image encoding: {e}")
+        raise
 
 @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)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
     fallback_image = Image.new("RGB", (width, height), (255, 0, 0))  # Red image as a fallback
 
-    if init_image is None:
-        # text2img case
-        try:
+    try:
+        if init_image is None:
+            # text2img case
             result = pipe(
                 prompt=prompt,
                 height=height,
@@ -62,53 +70,34 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
                 guidance_scale=0.0,
                 max_sequence_length=256
             )
-            image = result.images[0]
-            return image, seed
-        except Exception as e:
-            print(f"Pipeline call failed with error: {e}")
-            return fallback_image, seed
-    else:
-        # img2img case
-        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)
-        
-        print("Starting encoding of the image...")
-        latents = encode_image(init_image, vae)
+        else:
+            # img2img case
+            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)
+            
+            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), mode='bilinear')
+            
+            latent_channels = latents.shape[1]
+            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=flux_dtype)
+                latents = conv(latents.to(flux_dtype))
+
+            latents = latents.permute(0, 2, 3, 1).contiguous().view(-1, pipe.vae.config.latent_channels)
+            print(f"Latents shape after permutation: {latents.shape}")
 
-        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 = 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, 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)
-                _ = pipe.transformer(latents, timesteps=timestep)
-            else:
-                _ = pipe.transformer(latents)
-        except Exception as e:
-            print(f"Transformer call failed with error: {e}. Skipping transformer step.")
-            return fallback_image, seed
-
-        try:
-            print("Generating final image with the FLUX pipeline...")
-            image = pipe(
+            result = pipe(
                 prompt=prompt,
                 height=height,
                 width=width,
@@ -116,13 +105,13 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
                 generator=generator,
                 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
-    
-    return image, seed
+            )
+        
+        image = result.images[0]
+        return image, seed
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        return fallback_image, seed
 
 # Define example prompts
 examples = [