Update app.py

app.py

@@ -13,7 +13,6 @@ from diffusers.utils import load_image
 from PIL import Image
 
 import torch
-from torchvision import transforms
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -43,90 +42,325 @@ def init_predictor():
             compiler_transformer=False,
         )
     )
-    
-@spaces.GPU(duration=120)
-def generate_video(prompt, image, size, steps, frames, guidance_scale, progress=gr.Progress(track_tqdm=True) ):
-    print(f"image:{type(image)}")
+
+@spaces.GPU(duration=60)
+def generate_video(segment, image, prompt, size, guidance_scale, num_inference_steps, frames, seed, progress=gr.Progress(track_tqdm=True) ):
+
     random.seed(time.time())
     seed = int(random.randrange(4294967294))
-    img = load_image(image=image)
-    img.resize((size,size), Image.LANCZOS)
-    kwargs = {
-        "image": img,
-        "prompt": prompt,
-        "height": size,
-        "width": size,
-        "num_frames": frames,
-        "num_inference_steps": steps,
-        "seed": seed,
-        "guidance_scale": guidance_scale,
-        "embedded_guidance_scale": 1.0,
-        "negative_prompt": "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion",
-        "cfg_for": False,
-    }
-    assert image is not None, "please input image"
-
-    '''
-    preprocess = transforms.Compose(
-        [
-            transforms.ToTensor(), # Converts PIL [0, 255] (H, W, C) to Tensor [0, 1] (C, H, W)
-            transforms.Normalize([0.5], [0.5]), # Normalizes Tensor [0, 1] to [-1, 1]
-                                                 # Use [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] for RGB
-        ]
-    )
-    image_tensor = preprocess(img)
-    # 3. Add the batch dimension (B=1)
-    #    Resulting shape: [1, C, H, W]
-    image_tensor = image_tensor.unsqueeze(0)
-    kwargs["image"] = image_tensor
-    '''
+    if segment==1:
+
+
+        prompt_embeds, pooled_prompt_embeds, prompt_attention_mask = pipe.encode_prompt(
+        prompt=prompt, prompt_2=prompt, device=device
+        )
+        transformer_pooled_projections = pooled_prompt_embeds
+        transformer_pooled_projections = torch.cat([negative_pooled_prompt_embeds, pooled_prompt_embeds])
+        pipe.scheduler.set_timesteps(num_inference_steps, device=torch.device('cuda'))
+        timesteps = pipe.scheduler.timesteps
+        all_timesteps_cpu = timesteps.cpu()
+        timesteps_split_np = np.array_split(all_timesteps_cpu.numpy(), 8)
+        segment_timesteps = torch.from_numpy(timesteps_split_np[0]).to("cuda")
+
+        num_channels_latents = pipe.transformer.config.in_channels
+        latents = pipe.prepare_latents(
+            batch_size=1, num_channels_latents=pipe.transformer.config.in_channels, height=height, width=width, num_frames=frames,
+            dtype=torch.float32, device=device, generator=generator, latents=None,
+        )
+        guidance = torch.tensor([guidance_scale] * latents.shape[0], dtype=transformer_dtype, device=device) * 1000.0
+
+        kwargs = {
+            "prompt": prompt,
+            "height": size,
+            "width": size,
+            "num_frames": frames,
+            "num_inference_steps": steps,
+            "seed": seed,
+            "guidance_scale": guidance_scale,
+            "embedded_guidance_scale": 1.0,
+            "negative_prompt": "Aerial view, aerial view, overexposed, low quality, deformation, a poor composition, bad hands, bad teeth, bad eyes, bad limbs, distortion",
+            "cfg_for": False,
+        }
+        assert image is not None, "please input image"
+        img = load_image(image=image)
+        img.resize((size,size), Image.LANCZOS)
+        kwargs["image"] = img
+    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)
+
+        
+                current_latents = latents
+
+    for i, t in enumerate(pipe.progress_bar(segment_timesteps)):
+
+                latent_model_input = latents.to(transformer_dtype)
+        
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                with torch.no_grad():
+                  noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    timestep=timestep,
+                    encoder_hidden_states=prompt_embeds,
+                    encoder_attention_mask=prompt_attention_mask,
+                    pooled_projections=pooled_prompt_embeds,
+                    guidance=guidance,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=False,
+                  )[0]
+
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                else:
+                    video = latents
+                    return latents
+                    
+            intermediate_latents_cpu = current_latents.detach().cpu()
+
+    if segment==8:
+        latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
+        video = self.vae.decode(latents, return_dict=False)[0]
+        video = self.video_processor.postprocess_video(video, output_type=output_type)
+                    
+        return HunyuanVideoPipelineOutput(frames=video)
     
-    output = predictor.inference(kwargs)
-    save_dir = f"./"
-    video_out_file = f"{save_dir}/{seed}.mp4"
-    print(f"generate video, local path: {video_out_file}")
-    export_to_video(output, video_out_file, fps=24)
-    return video_out_file
+        save_dir = f"./"
+        video_out_file = f"{save_dir}/{seed}.mp4"
+        print(f"generate video, local path: {video_out_file}")
+        export_to_video(output, video_out_file, fps=24)
+    
+        return video_out_file, seed
+    else:
+        original_prompt_embeds_cpu = prompt_embeds.cpu()
+        original_negative_prompt_embeds_cpu = negative_prompt_embeds.cpu()
+        original_pooled_prompt_embeds_cpu = pooled_prompt_embeds.cpu()
+        original_negative_pooled_prompt_embeds_cpu = negative_pooled_prompt_embeds.cpu()
+        original_add_time_ids_cpu = add_time_ids.cpu()
+        timesteps = pipe.scheduler.timesteps
+        all_timesteps_cpu = timesteps.cpu() # Move to CPU
+        state = {
+            "intermediate_latents": intermediate_latents_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,
+            "pooled_prompt_embeds": original_pooled_prompt_embeds_cpu,
+            "negative_pooled_prompt_embeds": original_negative_pooled_prompt_embeds_cpu,
+            "add_time_ids": original_add_time_ids_cpu, # Save ORIGINAL time IDs
+            "guidance_scale": guidance_scale,
+            "timesteps_split": timesteps_split_for_state,
+            "seed": seed,
+            "prompt": prompt, # Save originals for reference/verification
+            "negative_prompt": negative_prompt,
+            "height": height, # Save dimensions used
+            "width": width
+        }
+        state_file = f"SkyReel_{segment}_{seed}.pt"
+        torch.save(state, state_file)
+        return None, seed
+
 
+    
 with gr.Blocks() as demo:
-            with gr.Row():
-                image = gr.Image(label="Upload Image", type="filepath")
-                prompt = gr.Textbox(label="Input Prompt")
-                size = gr.Slider(
-                    label="Size",
-                    minimum=256,
-                    maximum=1024,
-                    step=16,
-                    value=368,
-                )
-                frames = gr.Slider(
-                    label="Number of Frames",
-                    minimum=16,
-                    maximum=256,
-                    step=8,
-                    value=64,
-                )
-                steps = gr.Slider(
-                    label="Number of Steps",
-                    minimum=1,
-                    maximum=96,
-                    step=1,
-                    value=25,
-                )
-                guidance_scale = gr.Slider(
-                    label="Guidance Scale",
-                    minimum=1.0,
-                    maximum=16.0,
-                    step=.1,
-                    value=6.0,
-                )
-            submit_button = gr.Button("Generate Video")
-            output_video = gr.Video(label="Generated Video")
-            submit_button.click(
-                fn=generate_video,
-                inputs=[prompt, image, size, steps, frames, guidance_scale],
-                outputs=[output_video],
+        with gr.Row():
+            image = gr.Image(label="Upload Image", type="filepath")
+            prompt = gr.Textbox(label="Input Prompt")
+            size = gr.Slider(
+                label="Size",
+                minimum=256,
+                maximum=1024,
+                step=16,
+                value=368,
+            )
+            frames = gr.Slider(
+                label="Number of Frames",
+                minimum=16,
+                maximum=256,
+                step=8,
+                value=64,
+            )
+            steps = gr.Slider(
+                label="Number of Steps",
+                minimum=1,
+                maximum=96,
+                step=1,
+                value=25,
+            )
+            guidance_scale = gr.Slider(
+                label="Guidance Scale",
+                minimum=1.0,
+                maximum=16.0,
+                step=.1,
+                value=6.0,
+            )
+        submit_button = gr.Button("Generate Video")
+        output_video = gr.Video(label="Generated Video")
+        range_sliders = []
+        for i in range(8):
+            slider = gr.Slider(
+                minimum=1,
+                maximum=250,
+                value=[i * (num_inference_steps.value // 8)],
+                step=1,
+                label=f"Range {i + 1}",
             )
+            range_sliders.append(slider)
+        num_inference_steps.change(
+        update_ranges,
+        inputs=num_inference_steps,
+        outputs=range_sliders,
+        )
+        
+        gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        cache_examples=False
+    )
+    use_negative_prompt.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt,
+        outputs=negative_prompt,
+        api_name=False,
+    )
+    gr.on(
+        triggers=[
+            run_button_1.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    gr.on(
+        triggers=[
+            run_button_2.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    gr.on(
+        triggers=[
+            run_button_3.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    gr.on(
+        triggers=[
+            run_button_4.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    gr.on(
+        triggers=[
+            run_button_5.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    gr.on(
+        triggers=[
+            run_button_6.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    gr.on(
+        triggers=[
+            run_button_7.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    gr.on(
+        triggers=[
+            run_button_8.click,
+        ],
+        fn=generate,
+        inputs=[
+            gr.Number(value=4),
+            image,
+            prompt,
+            size,
+            guidance_scale,
+            num_inference_steps,
+            frames,
+            seed,
+        ],
+        outputs=[result, seed],
+    )
+    
     
 if __name__ == "__main__":
     init_predictor()