Update app.py

app.py

@@ -110,16 +110,16 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         segment_timesteps = torch.from_numpy(timesteps_split_np[0]).to("cuda")
         num_channels_latents = pipe.transformer.config.in_channels
         num_channels_latents = int(num_channels_latents / 2)
-        image = pipe.video_processor.preprocess(image, height=height, width=width).to(
+        image = pipe.video_processor.preprocess(image, height=size, width=size).to(
                 device, dtype=prompt_embeds.dtype
         )
         num_latent_frames = (frames - 1) // pipe.vae_scale_factor_temporal + 1
         latents = pipe.prepare_latents(
-            batch_size=1, num_channels_latents=pipe.transformer.config.in_channels, height=height, width=width, num_frames=frames,
+            batch_size=1, num_channels_latents=pipe.transformer.config.in_channels, height=size, width=size, num_frames=frames,
             dtype=torch.float32, device=device, generator=generator, latents=None,
         )
         image_latents = pipe.image_latents(
-                image, batch_size, height, width, device, torch.float32, num_channels_latents, num_latent_frames
+                image, batch_size, size, size, device, torch.float32, num_channels_latents, num_latent_frames
         )
         image_latents = image_latents.to(pipe.transformer.dtype)
         guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
@@ -130,8 +130,8 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
         latents = state["intermediate_latents"].to("cuda", dtype=torch.bfloat16)
         guidance_scale = state["guidance_scale"]
         all_timesteps_cpu = state["all_timesteps"]
-        height = state["height"]
-        width = state["width"]
+        size = state["height"]
+        size = state["width"]
         pipe.scheduler.set_timesteps(len(all_timesteps_cpu), device=device)
         timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
         segment_timesteps = torch.from_numpy(timesteps_split_np[segment - 1]).to("cuda")
@@ -194,8 +194,8 @@ def generate(segment, image, prompt, size, guidance_scale, num_inference_steps,
             "seed": seed,
             "prompt": prompt, # Save originals for reference/verification
             "negative_prompt": negative_prompt,
-            "height": height, # Save dimensions used
-            "width": width
+            "height": size, # Save dimensions used
+            "width": size
         }
         state_file = f"SkyReel_{segment}_{seed}.pt"
         torch.save(state, state_file)