soft fail

app.py

@@ -62,10 +62,6 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             image = result.images[0]
             print(f"Generated image shape: {image.size}")
 
-            # Inspect the output and log relevant details
-            print("Logging detailed information for text2img:")
-            # Log intermediate latent information if accessible
-
             print("Logging complete.")
         except Exception as e:
             print(f"Pipeline call failed with error: {e}")
@@ -95,13 +91,16 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
 
         try:
             print("Calling the transformer with latents")
-            # Initialize timestep variable
-            timestep = torch.tensor([num_inference_steps], device=device, dtype=dtype)
-            _ = pipe.transformer(latents, timestep=timestep)
+            # 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}")
-            raise
+            print(f"Transformer call failed with error: {e}. Skipping transformer step.")
+            return "Transformer call failed, skipping the step."
 
         try:
             print("Calling the diffusion pipeline with latents")
@@ -116,11 +115,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}")
-            raise
+            return f"Pipeline call with latents failed: {e}"
     
     print("Inference complete")
     return image, seed
 
+
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",