fix: out of memory error

app.py

@@ -10,6 +10,8 @@ if gr.NO_RELOAD:
     high_noise_frac = 0.8
     negative_prompt = "Distorted, discontinuous, Ugly, blurry, low resolution, motionless, static, disfigured, disconnected limbs, Ugly faces, incomplete arms"
     generator = torch.manual_seed(8888)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("Device:", device)
 
     base = DiffusionPipeline.from_pretrained(
         "stabilityai/stable-diffusion-xl-base-1.0", 
@@ -35,18 +37,21 @@ if gr.NO_RELOAD:
     refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)
     pipeline.unet = torch.compile(pipeline.unet, mode="reduce-overhead", fullgraph=True)
 
-def generate(prompt: str):
+def generate(prompt: str, progress=gr.Progress()):
+    progress((0, 100), desc="Starting..")
     image = base(
         prompt=prompt,
         num_inference_steps=n_steps,
         denoising_end=high_noise_frac,
         output_type="latent",
+        callback=lambda s, t, d: progress((s, 100), desc="Generating first frame..."),
     ).images[0]
     image = refiner(
         prompt=prompt,
         num_inference_steps=n_steps,
         denoising_start=high_noise_frac,
         image=image,
+        callback=lambda s, t, d: progress((s+40, 100), desc="Refining first frame..."),
     ).images[0]
     image = to_tensor(image)
     frames: list[Image.Image] = pipeline(
@@ -56,6 +61,8 @@ def generate(prompt: str):
         negative_prompt=negative_prompt,
         guidance_scale=9.0,
         generator=generator,
+        decode_chunk_size=10,
+        callback=lambda s, t, d: progress((s+50, 100), desc="Generating video..."),
     ).frames[0]
     frames = [to_tensor(frame.convert("RGB")).mul(255).byte().permute(1, 2, 0) for frame in frames]
     frames = torch.stack(frames)