Update app.py

app.py

@@ -29,7 +29,7 @@ pokemon_model.decoder.compile(jit_compile=True)
 pokemon_model.text_encoder.compile(jit_compile=True)
 
 # Warm-up the model.
-_ = pokemon_model.text_to_image("Teddy bear", batch_size=num_images_to_gen)
+#_ = pokemon_model.text_to_image("Teddy bear", batch_size=num_images_to_gen)
 
 
 def generate_image_fn(prompt: str, unconditional_guidance_scale: int) -> list:
@@ -37,6 +37,7 @@ def generate_image_fn(prompt: str, unconditional_guidance_scale: int) -> list:
     # `images is an `np.ndarray`. So we convert it to a list of ndarrays.
     # Each ndarray represents a generated image.
     # Reference: https://gradio.app/docs/#gallery
+    pokemon_model.to("cpu")
     images = pokemon_model.text_to_image(
         prompt,
         batch_size=num_images_to_gen,