handling

app.py

@@ -13,43 +13,37 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Load the diffusion pipeline
-pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)
+# Load the 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)
 
 def preprocess_image(image, image_size):
-    print(f"Preprocessing image to size: {image_size}x{image_size}")
     preprocess = transforms.Compose([
-        transforms.Resize((image_size, image_size)),  # Use model-specific size
+        transforms.Resize((image_size, image_size)),
         transforms.ToTensor(),
-        transforms.Normalize([0.5], [0.5])  # Ensure this matches the VAE's training normalization
+        transforms.Normalize([0.5], [0.5])
     ])
     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")
     with torch.no_grad():
         latents = vae.encode(image).latent_dist.sample() * 0.18215
-    print(f"Latents shape after encoding: {latents.shape}")
     return latents
 
-# A utility function to log shapes and other relevant information
-def log_tensor_info(tensor, name):
-    print(f"{name} shape: {tensor.shape} dtype: {tensor.dtype} device: {tensor.device}")
-
 @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)
 
+    fallback_image = Image.new("RGB", (width, height), (255, 0, 0))  # Red image as a fallback
+
     if init_image is None:
-        print("No initial image provided, processing text2img")
         try:
-            print("Calling the diffusion pipeline for text2img")
             result = pipe(
                 prompt=prompt,
                 height=height,
@@ -60,50 +54,35 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
                 max_sequence_length=256
             )
             image = result.images[0]
-            print(f"Generated image shape: {image.size}")
-
-            print("Logging complete.")
+            return image, seed
         except Exception as e:
             print(f"Pipeline call failed with error: {e}")
-            raise
+            return fallback_image, seed
     else:
-        print("Initial image provided, processing img2img")
-        vae_image_size = pipe.vae.config.sample_size
-        print(f"Expected VAE image size: {vae_image_size}")
+        vae_image_size = pipe.vae.config.sample_size  # Ensure this is correct
         init_image = init_image.convert("RGB")
         init_image = preprocess_image(init_image, vae_image_size)
         latents = encode_image(init_image, pipe.vae)
 
-        print("Interpolating latents to match model's input size...")
         latents = torch.nn.functional.interpolate(latents, size=(height // 8, width // 8))
-        log_tensor_info(latents, "Latents after interpolation")
-
-        latent_channels = pipe.vae.config.latent_channels
-        print(f"Expected latent channels: 64, current latent channels: {latent_channels}")
+        latent_channels = pipe.vae.config.latent_channels  # Ensure this is correct
         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)
-            log_tensor_info(latents, "Latents after channel conversion")
 
         latents = latents.permute(0, 2, 3, 1).contiguous().view(-1, 64)
-        log_tensor_info(latents, "Latents after reshaping for transformer")
 
         try:
-            print("Calling the transformer with latents")
-            # Check if timestep is required and initialize it if necessary
             if '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)
-            print("Transformer call succeeded")
         except Exception as e:
             print(f"Transformer call failed with error: {e}. Skipping transformer step.")
-            return "Transformer call failed, skipping the step."
+            return fallback_image, seed
 
         try:
-            print("Calling the diffusion pipeline with latents")
             image = pipe(
                 prompt=prompt,
                 height=height,
@@ -115,12 +94,12 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             ).images[0]
         except Exception as e:
             print(f"Pipeline call with latents failed with error: {e}")
-            return f"Pipeline call with latents failed: {e}"
+            return fallback_image, seed
     
-    print("Inference complete")
     return image, seed
 
 
+
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",