safe check

app.py

@@ -43,6 +43,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
     fallback_image = Image.new("RGB", (width, height), (255, 0, 0))  # Red image as a fallback
 
     if init_image is None:
+        # text2img case
         try:
             result = pipe(
                 prompt=prompt,
@@ -59,6 +60,7 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
             print(f"Pipeline call failed with error: {e}")
             return fallback_image, seed
     else:
+        # img2img case
         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)
@@ -73,7 +75,8 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
         latents = latents.permute(0, 2, 3, 1).contiguous().view(-1, 64)
 
         try:
-            if 'timesteps' in pipe.transformer.forward.__code__.co_varnames:
+            # 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:
@@ -98,8 +101,6 @@ def infer(prompt, init_image=None, seed=42, randomize_seed=False, width=1024, he
     
     return image, seed
 
-
-
 # Define example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",