Update app.py

app.py

@@ -1,9 +1,7 @@
 import os
-
 os.environ['HF_HOME'] = os.path.abspath(
     os.path.realpath(os.path.join(os.path.dirname(__file__), './hf_download'))
 )
-
 import gradio as gr
 import torch
 import traceback
@@ -12,7 +10,6 @@ import safetensors.torch as sf
 import numpy as np
 import math
 import spaces
-
 from PIL import Image
 from diffusers import AutoencoderKLHunyuanVideo
 from transformers import (
@@ -31,68 +28,54 @@ from diffusers_helper.utils import (
 )
 from diffusers_helper.models.hunyuan_video_packed import HunyuanVideoTransformer3DModelPacked
 from diffusers_helper.pipelines.k_diffusion_hunyuan import sample_hunyuan
-from diffusers_helper.memory import (
-    cpu, gpu,
-    get_cuda_free_memory_gb,
-    move_model_to_device_with_memory_preservation,
-    offload_model_from_device_for_memory_preservation,
-    fake_diffusers_current_device,
-    DynamicSwapInstaller,
-    unload_complete_models,
-    load_model_as_complete
-)
-from diffusers_helper.thread_utils import AsyncStream, async_run
-from diffusers_helper.gradio.progress_bar import make_progress_bar_css, make_progress_bar_html
-from transformers import SiglipImageProcessor, SiglipVisionModel
-from diffusers_helper.clip_vision import hf_clip_vision_encode
-from diffusers_helper.bucket_tools import find_nearest_bucket
 
-# Check GPU memory
-free_mem_gb = get_cuda_free_memory_gb(gpu)
-high_vram = free_mem_gb > 60
-
-print(f'Free VRAM {free_mem_gb} GB')
-print(f'High-VRAM Mode: {high_vram}')
+# Remove or replace GPU-specific imports
+device = torch.device("cpu")
 
 # Load models
 text_encoder = LlamaModel.from_pretrained(
     "hunyuanvideo-community/HunyuanVideo",
     subfolder='text_encoder',
     torch_dtype=torch.float16
-).cpu()
+).to(device)
+
 text_encoder_2 = CLIPTextModel.from_pretrained(
     "hunyuanvideo-community/HunyuanVideo",
     subfolder='text_encoder_2',
     torch_dtype=torch.float16
-).cpu()
+).to(device)
+
 tokenizer = LlamaTokenizerFast.from_pretrained(
     "hunyuanvideo-community/HunyuanVideo",
     subfolder='tokenizer'
 )
+
 tokenizer_2 = CLIPTokenizer.from_pretrained(
     "hunyuanvideo-community/HunyuanVideo",
     subfolder='tokenizer_2'
 )
+
 vae = AutoencoderKLHunyuanVideo.from_pretrained(
     "hunyuanvideo-community/HunyuanVideo",
     subfolder='vae',
     torch_dtype=torch.float16
-).cpu()
+).to(device)
 
 feature_extractor = SiglipImageProcessor.from_pretrained(
     "lllyasviel/flux_redux_bfl",
     subfolder='feature_extractor'
 )
+
 image_encoder = SiglipVisionModel.from_pretrained(
     "lllyasviel/flux_redux_bfl",
     subfolder='image_encoder',
     torch_dtype=torch.float16
-).cpu()
+).to(device)
 
 transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained(
     'lllyasviel/FramePack_F1_I2V_HY_20250503',
     torch_dtype=torch.bfloat16
-).cpu()
+).to(device)
 
 # Evaluation mode
 vae.eval()
@@ -101,14 +84,6 @@ text_encoder_2.eval()
 image_encoder.eval()
 transformer.eval()
 
-# Slicing/Tiling for low VRAM
-if not high_vram:
-    vae.enable_slicing()
-    vae.enable_tiling()
-
-transformer.high_quality_fp32_output_for_inference = True
-print('transformer.high_quality_fp32_output_for_inference = True')
-
 # Move to correct dtype
 transformer.to(dtype=torch.bfloat16)
 vae.to(dtype=torch.float16)
@@ -123,19 +98,7 @@ text_encoder_2.requires_grad_(False)
 image_encoder.requires_grad_(False)
 transformer.requires_grad_(False)
 
-# DynamicSwap if low VRAM
-if not high_vram:
-    DynamicSwapInstaller.install_model(transformer, device=gpu)
-    DynamicSwapInstaller.install_model(text_encoder, device=gpu)
-else:
-    text_encoder.to(gpu)
-    text_encoder_2.to(gpu)
-    image_encoder.to(gpu)
-    vae.to(gpu)
-    transformer.to(gpu)
-
 stream = AsyncStream()
-
 outputs_folder = './outputs/'
 os.makedirs(outputs_folder, exist_ok=True)
 
@@ -145,7 +108,6 @@ examples = [
     ["img_examples/3.png", "The woman dances elegantly among the blossoms, spinning slowly with flowing sleeves and graceful hand movements."]
 ]
 
-# Example generation (optional)
 def generate_examples(input_image, prompt):
     t2v=False
     n_prompt=""
@@ -156,32 +118,24 @@ def generate_examples(input_image, prompt):
     cfg=1.0
     gs=10.0
     rs=0.0
-    gpu_memory_preservation=6
+    gpu_memory_preservation=6  # unused
     use_teacache=True
     mp4_crf=16
-
     global stream
-
     if t2v:
         default_height, default_width = 640, 640
         input_image = np.ones((default_height, default_width, 3), dtype=np.uint8) * 255
         print("No input image provided. Using a blank white image.")
-
     yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
-
     stream = AsyncStream()
-
     async_run(
         worker, input_image, prompt, n_prompt, seed,
         total_second_length, latent_window_size, steps,
         cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
     )
-
     output_filename = None
-
     while True:
         flag, data = stream.output_queue.next()
-
         if flag == 'file':
             output_filename = data
             yield (
@@ -192,7 +146,6 @@ def generate_examples(input_image, prompt):
                 gr.update(interactive=False),
                 gr.update(interactive=True)
             )
-
         if flag == 'progress':
             preview, desc, html = data
             yield (
@@ -203,7 +156,6 @@ def generate_examples(input_image, prompt):
                 gr.update(interactive=False),
                 gr.update(interactive=True)
             )
-
         if flag == 'end':
             yield (
                 output_filename,
@@ -221,84 +173,44 @@ def worker(
     total_second_length, latent_window_size, steps,
     cfg, gs, rs, gpu_memory_preservation, use_teacache, mp4_crf
 ):
-    # Calculate total sections
     total_latent_sections = (total_second_length * 30) / (latent_window_size * 4)
     total_latent_sections = int(max(round(total_latent_sections), 1))
-
     job_id = generate_timestamp()
-
     stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Starting ...'))))
 
     try:
-        # Unload if VRAM is low
-        if not high_vram:
-            unload_complete_models(
-                text_encoder, text_encoder_2, image_encoder, vae, transformer
-            )
-
-        # Text encoding
-        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Text encoding ...'))))
-
-        if not high_vram:
-            fake_diffusers_current_device(text_encoder, gpu)
-            load_model_as_complete(text_encoder_2, target_device=gpu)
-
         llama_vec, clip_l_pooler = encode_prompt_conds(prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
-
         if cfg == 1:
             llama_vec_n, clip_l_pooler_n = torch.zeros_like(llama_vec), torch.zeros_like(clip_l_pooler)
         else:
             llama_vec_n, clip_l_pooler_n = encode_prompt_conds(n_prompt, text_encoder, text_encoder_2, tokenizer, tokenizer_2)
-
         llama_vec, llama_attention_mask = crop_or_pad_yield_mask(llama_vec, length=512)
         llama_vec_n, llama_attention_mask_n = crop_or_pad_yield_mask(llama_vec_n, length=512)
 
-        # Process image
-        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Image processing ...'))))
-
         H, W, C = input_image.shape
         height, width = find_nearest_bucket(H, W, resolution=640)
         input_image_np = resize_and_center_crop(input_image, target_width=width, target_height=height)
-
         Image.fromarray(input_image_np).save(os.path.join(outputs_folder, f'{job_id}.png'))
-
         input_image_pt = torch.from_numpy(input_image_np).float() / 127.5 - 1
         input_image_pt = input_image_pt.permute(2, 0, 1)[None, :, None]
 
-        # VAE encoding
-        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'VAE encoding ...'))))
-
-        if not high_vram:
-            load_model_as_complete(vae, target_device=gpu)
-        start_latent = vae_encode(input_image_pt, vae)
-
-        # CLIP Vision
-        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'CLIP Vision encoding ...'))))
+        start_latent = vae_encode(input_image_pt, vae).to(device)
 
-        if not high_vram:
-            load_model_as_complete(image_encoder, target_device=gpu)
         image_encoder_output = hf_clip_vision_encode(input_image_np, feature_extractor, image_encoder)
         image_encoder_last_hidden_state = image_encoder_output.last_hidden_state
 
-        # Convert dtype
-        llama_vec = llama_vec.to(transformer.dtype)
-        llama_vec_n = llama_vec_n.to(transformer.dtype)
-        clip_l_pooler = clip_l_pooler.to(transformer.dtype)
-        clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype)
-        image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype)
-
-        # Start sampling
-        stream.output_queue.push(('progress', (None, '', make_progress_bar_html(0, 'Start sampling ...'))))
+        llama_vec = llama_vec.to(transformer.dtype).to(device)
+        llama_vec_n = llama_vec_n.to(transformer.dtype).to(device)
+        clip_l_pooler = clip_l_pooler.to(transformer.dtype).to(device)
+        clip_l_pooler_n = clip_l_pooler_n.to(transformer.dtype).to(device)
+        image_encoder_last_hidden_state = image_encoder_last_hidden_state.to(transformer.dtype).to(device)
 
         rnd = torch.Generator("cpu").manual_seed(seed)
-
         history_latents = torch.zeros(
             size=(1, 16, 16 + 2 + 1, height // 8, width // 8),
             dtype=torch.float32
-        ).cpu()
-        history_pixels = None
+        ).to(device)
 
-        # Add start_latent
         history_latents = torch.cat([history_latents, start_latent.to(history_latents)], dim=2)
         total_generated_latent_frames = 1
 
@@ -307,15 +219,6 @@ def worker(
                 stream.output_queue.push(('end', None))
                 return
 
-            print(f'section_index = {section_index}, total_latent_sections = {total_latent_sections}')
-
-            if not high_vram:
-                unload_complete_models()
-                move_model_to_device_with_memory_preservation(
-                    transformer, target_device=gpu,
-                    preserved_memory_gb=gpu_memory_preservation
-                )
-
             if use_teacache:
                 transformer.initialize_teacache(enable_teacache=True, num_steps=steps)
             else:
@@ -326,11 +229,9 @@ def worker(
                 preview = vae_decode_fake(preview)
                 preview = (preview * 255.0).detach().cpu().numpy().clip(0, 255).astype(np.uint8)
                 preview = einops.rearrange(preview, 'b c t h w -> (b h) (t w) c')
-
                 if stream.input_queue.top() == 'end':
                     stream.output_queue.push(('end', None))
                     raise KeyboardInterrupt('User ends the task.')
-
                 current_step = d['i'] + 1
                 percentage = int(100.0 * current_step / steps)
                 hint = f'Sampling {current_step}/{steps}'
@@ -350,11 +251,9 @@ def worker(
             ) = indices.split([1, 16, 2, 1, latent_window_size], dim=1)
 
             clean_latent_indices = torch.cat([clean_latent_indices_start, clean_latent_1x_indices], dim=1)
-
             clean_latents_4x, clean_latents_2x, clean_latents_1x = history_latents[
                 :, :, -sum([16, 2, 1]):, :, :
             ].split([16, 2, 1], dim=2)
-
             clean_latents = torch.cat(
                 [start_latent.to(history_latents), clean_latents_1x],
                 dim=2
@@ -377,7 +276,7 @@ def worker(
                 negative_prompt_embeds=llama_vec_n,
                 negative_prompt_embeds_mask=llama_attention_mask_n,
                 negative_prompt_poolers=clip_l_pooler_n,
-                device=gpu,
+                device=device,
                 dtype=torch.bfloat16,
                 image_embeddings=image_encoder_last_hidden_state,
                 latent_indices=latent_indices,
@@ -393,18 +292,12 @@ def worker(
             total_generated_latent_frames += int(generated_latents.shape[2])
             history_latents = torch.cat([history_latents, generated_latents.to(history_latents)], dim=2)
 
-            if not high_vram:
-                offload_model_from_device_for_memory_preservation(transformer, target_device=gpu, preserved_memory_gb=8)
-                load_model_as_complete(vae, target_device=gpu)
-
             real_history_latents = history_latents[:, :, -total_generated_latent_frames:, :, :]
-
             if history_pixels is None:
                 history_pixels = vae_decode(real_history_latents, vae).cpu()
             else:
                 section_latent_frames = latent_window_size * 2
                 overlapped_frames = latent_window_size * 4 - 3
-
                 current_pixels = vae_decode(
                     real_history_latents[:, :, -section_latent_frames:], vae
                 ).cpu()
@@ -412,21 +305,12 @@ def worker(
                     history_pixels, current_pixels, overlapped_frames
                 )
 
-            if not high_vram:
-                unload_complete_models()
-
             output_filename = os.path.join(outputs_folder, f'{job_id}_{total_generated_latent_frames}.mp4')
-
             save_bcthw_as_mp4(history_pixels, output_filename, fps=30)
-
-            print(f'Decoded. Latent shape {real_history_latents.shape}; pixel shape {history_pixels.shape}')
-
             stream.output_queue.push(('file', output_filename))
 
-    except:
+    except Exception as e:
         traceback.print_exc()
-        if not high_vram:
-            unload_complete_models(text_encoder, text_encoder_2, image_encoder, vae, transformer)
 
     stream.output_queue.push(('end', None))
     return