Update app.py

app.py

@@ -183,7 +183,9 @@ def generate(
     else:
         state_file = f"rv_L_{segment-1}_{seed}.pt"
         state = torch.load(state_file, weights_only=False)
-        generator = torch.Generator(device='cuda').manual_seed(seed)
+        generator_state = state["generator_state"]
+        generator = torch.Generator(device='cuda') #.manual_seed(seed)
+        generator.set_state(generator_state)
         latents = state["intermediate_latents"].to("cuda", dtype=torch.bfloat16)
         guidance_scale = state["guidance_scale"]
         all_timesteps_cpu = state["all_timesteps"]
@@ -248,9 +250,11 @@ def generate(
         original_pooled_prompt_embeds_cpu = pooled_prompt_embeds.cpu()
         original_negative_pooled_prompt_embeds_cpu = negative_pooled_prompt_embeds.cpu()
         timesteps = pipe.scheduler.timesteps
+        generator_state_cpu = generator_state.cpu()
         all_timesteps_cpu = timesteps.cpu() # Move to CPU
         state = {
         "intermediate_latents": intermediate_latents_cpu,
+        "generator_state": generator_state_cpu,
         "all_timesteps": all_timesteps_cpu, # Save full list generated by scheduler
         "prompt_embeds": original_prompt_embeds_cpu, # Save ORIGINAL embeds
         "negative_prompt_embeds": original_negative_prompt_embeds_cpu,