修复Hugging Face Space GPU支持问题

app.py

@@ -12,6 +12,17 @@ import safetensors.torch as sf
 import numpy as np
 import math
 
+# 检查是否在Hugging Face Space环境中
+IN_HF_SPACE = os.environ.get('SPACE_ID') is not None
+
+# 如果在Hugging Face Space中，导入spaces模块
+if IN_HF_SPACE:
+    try:
+        import spaces
+        print("在Hugging Face Space环境中运行，已导入spaces模块")
+    except ImportError:
+        print("未能导入spaces模块，可能不在Hugging Face Space环境中")
+
 from PIL import Image
 from diffusers import AutoencoderKLHunyuanVideo
 from transformers import LlamaModel, CLIPTextModel, LlamaTokenizerFast, CLIPTokenizer
@@ -27,59 +38,86 @@ from diffusers_helper.clip_vision import hf_clip_vision_encode
 from diffusers_helper.bucket_tools import find_nearest_bucket
 
 # 获取可用的CUDA内存
-free_mem_gb = get_cuda_free_memory_gb(gpu)
-high_vram = free_mem_gb > 60
+try:
+    if torch.cuda.is_available():
+        free_mem_gb = get_cuda_free_memory_gb(gpu)
+        print(f'Free VRAM {free_mem_gb} GB')
+    else:
+        free_mem_gb = 6.0  # 默认值
+        print("CUDA不可用，使用默认的内存设置")
+except Exception as e:
+    free_mem_gb = 6.0  # 默认值
+    print(f"获取CUDA内存时出错: {e}，使用默认的内存设置")
 
-print(f'Free VRAM {free_mem_gb} GB')
+high_vram = free_mem_gb > 60
 print(f'High-VRAM Mode: {high_vram}')
 
-# 加载模型
-text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
-text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
-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()
-
-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()
-
-transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePackI2V_HY', torch_dtype=torch.bfloat16).cpu()
-
-vae.eval()
-text_encoder.eval()
-text_encoder_2.eval()
-image_encoder.eval()
-transformer.eval()
-
-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')
-
-transformer.to(dtype=torch.bfloat16)
-vae.to(dtype=torch.float16)
-image_encoder.to(dtype=torch.float16)
-text_encoder.to(dtype=torch.float16)
-text_encoder_2.to(dtype=torch.float16)
-
-vae.requires_grad_(False)
-text_encoder.requires_grad_(False)
-text_encoder_2.requires_grad_(False)
-image_encoder.requires_grad_(False)
-transformer.requires_grad_(False)
-
-if not high_vram:
-    # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
-    DynamicSwapInstaller.install_model(transformer, device=gpu)
-    DynamicSwapInstaller.install_model(text_encoder, device=gpu)
+# 使用加载模型的函数
+def load_models():
+    print("开始加载模型...")
+    
+    # 加载模型
+    text_encoder = LlamaModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder', torch_dtype=torch.float16).cpu()
+    text_encoder_2 = CLIPTextModel.from_pretrained("hunyuanvideo-community/HunyuanVideo", subfolder='text_encoder_2', torch_dtype=torch.float16).cpu()
+    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()
+
+    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()
+
+    transformer = HunyuanVideoTransformer3DModelPacked.from_pretrained('lllyasviel/FramePackI2V_HY', torch_dtype=torch.bfloat16).cpu()
+
+    vae.eval()
+    text_encoder.eval()
+    text_encoder_2.eval()
+    image_encoder.eval()
+    transformer.eval()
+
+    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')
+
+    transformer.to(dtype=torch.bfloat16)
+    vae.to(dtype=torch.float16)
+    image_encoder.to(dtype=torch.float16)
+    text_encoder.to(dtype=torch.float16)
+    text_encoder_2.to(dtype=torch.float16)
+
+    vae.requires_grad_(False)
+    text_encoder.requires_grad_(False)
+    text_encoder_2.requires_grad_(False)
+    image_encoder.requires_grad_(False)
+    transformer.requires_grad_(False)
+
+    if torch.cuda.is_available() and gpu.type == 'cuda':
+        if not high_vram:
+            # DynamicSwapInstaller is same as huggingface's enable_sequential_offload but 3x faster
+            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)
+    
+    return text_encoder, text_encoder_2, tokenizer, tokenizer_2, vae, feature_extractor, image_encoder, transformer
+
+# 使用Hugging Face Spaces GPU装饰器
+if IN_HF_SPACE and 'spaces' in globals():
+    @spaces.GPU
+    def load_models_with_gpu():
+        return load_models()
+    
+    print("使用@spaces.GPU装饰器加载模型")
+    text_encoder, text_encoder_2, tokenizer, tokenizer_2, vae, feature_extractor, image_encoder, transformer = load_models_with_gpu()
 else:
-    text_encoder.to(gpu)
-    text_encoder_2.to(gpu)
-    image_encoder.to(gpu)
-    vae.to(gpu)
-    transformer.to(gpu)
+    print("不使用@spaces.GPU装饰器，直接加载模型")
+    text_encoder, text_encoder_2, tokenizer, tokenizer_2, vae, feature_extractor, image_encoder, transformer = load_models()
 
 stream = AsyncStream()
 
@@ -303,32 +341,64 @@ def worker(input_image, prompt, n_prompt, seed, total_second_length, latent_wind
     return
 
 
-def process(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache):
-    global stream
-    assert input_image is not None, 'No input image!'
+# 使用Hugging Face Spaces GPU装饰器处理进程函数
+if IN_HF_SPACE and 'spaces' in globals():
+    @spaces.GPU
+    def process_with_gpu(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache):
+        global stream
+        assert input_image is not None, 'No input image!'
+
+        yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
+
+        stream = AsyncStream()
 
-    yield None, None, '', '', gr.update(interactive=False), gr.update(interactive=True)
+        async_run(worker, input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache)
 
-    stream = AsyncStream()
+        output_filename = None
 
-    async_run(worker, input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache)
+        while True:
+            flag, data = stream.output_queue.next()
 
-    output_filename = None
+            if flag == 'file':
+                output_filename = data
+                yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
 
-    while True:
-        flag, data = stream.output_queue.next()
+            if flag == 'progress':
+                preview, desc, html = data
+                yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
 
-        if flag == 'file':
-            output_filename = data
-            yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
+            if flag == 'end':
+                yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
+                break
+    
+    process = process_with_gpu
+else:
+    def process(input_image, prompt, n_prompt, seed, total_second_length, latent_window_size, steps, cfg, gs, rs, gpu_memory_preservation, use_teacache):
+        global stream
+        assert input_image is not None, 'No input 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)
 
-        if flag == 'progress':
-            preview, desc, html = data
-            yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
+        output_filename = None
 
-        if flag == 'end':
-            yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
-            break
+        while True:
+            flag, data = stream.output_queue.next()
+
+            if flag == 'file':
+                output_filename = data
+                yield output_filename, gr.update(), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=True)
+
+            if flag == 'progress':
+                preview, desc, html = data
+                yield gr.update(), gr.update(visible=True, value=preview), desc, html, gr.update(interactive=False), gr.update(interactive=True)
+
+            if flag == 'end':
+                yield output_filename, gr.update(visible=False), gr.update(), '', gr.update(interactive=True), gr.update(interactive=False)
+                break
 
 
 def end_process():

diffusers_helper/memory.py

@@ -2,10 +2,26 @@
 
 
 import torch
+import os
 
+# 检查是否在Hugging Face Space环境中
+IN_HF_SPACE = os.environ.get('SPACE_ID') is not None
 
+# 设置CPU设备
 cpu = torch.device('cpu')
-gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
+
+# 尝试设置GPU设备，如果不可用则回退到CPU
+try:
+    if torch.cuda.is_available():
+        gpu = torch.device(f'cuda:{torch.cuda.current_device()}')
+    else:
+        print("CUDA不可用，使用CPU作为默认设备")
+        gpu = torch.device('cpu')
+except Exception as e:
+    print(f"初始化CUDA设备时出错: {e}")
+    print("回退到CPU设备")
+    gpu = torch.device('cpu')
+
 gpu_complete_modules = []
 
 
@@ -71,19 +87,34 @@ def fake_diffusers_current_device(model: torch.nn.Module, target_device: torch.d
 def get_cuda_free_memory_gb(device=None):
     if device is None:
         device = gpu
-
-    memory_stats = torch.cuda.memory_stats(device)
-    bytes_active = memory_stats['active_bytes.all.current']
-    bytes_reserved = memory_stats['reserved_bytes.all.current']
-    bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
-    bytes_inactive_reserved = bytes_reserved - bytes_active
-    bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
-    return bytes_total_available / (1024 ** 3)
+    
+    # 如果不是CUDA设备，返回默认值
+    if device.type != 'cuda':
+        print("无法获取非CUDA设备的内存信息，返回默认值")
+        return 6.0  # 返回一个默认值
+    
+    try:
+        memory_stats = torch.cuda.memory_stats(device)
+        bytes_active = memory_stats['active_bytes.all.current']
+        bytes_reserved = memory_stats['reserved_bytes.all.current']
+        bytes_free_cuda, _ = torch.cuda.mem_get_info(device)
+        bytes_inactive_reserved = bytes_reserved - bytes_active
+        bytes_total_available = bytes_free_cuda + bytes_inactive_reserved
+        return bytes_total_available / (1024 ** 3)
+    except Exception as e:
+        print(f"获取CUDA内存信息时出错: {e}")
+        return 6.0  # 返回一个默认值
 
 
 def move_model_to_device_with_memory_preservation(model, target_device, preserved_memory_gb=0):
     print(f'Moving {model.__class__.__name__} to {target_device} with preserved memory: {preserved_memory_gb} GB')
 
+    # 如果目标设备是CPU或当前在CPU上，直接移动
+    if target_device.type == 'cpu' or gpu.type == 'cpu':
+        model.to(device=target_device)
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        return
+
     for m in model.modules():
         if get_cuda_free_memory_gb(target_device) <= preserved_memory_gb:
             torch.cuda.empty_cache()
@@ -100,6 +131,12 @@ def move_model_to_device_with_memory_preservation(model, target_device, preserve
 def offload_model_from_device_for_memory_preservation(model, target_device, preserved_memory_gb=0):
     print(f'Offloading {model.__class__.__name__} from {target_device} to preserve memory: {preserved_memory_gb} GB')
 
+    # 如果目标设备是CPU或当前在CPU上，直接处理
+    if target_device.type == 'cpu' or gpu.type == 'cpu':
+        model.to(device=cpu)
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
+        return
+
     for m in model.modules():
         if get_cuda_free_memory_gb(target_device) >= preserved_memory_gb:
             torch.cuda.empty_cache()
@@ -119,7 +156,7 @@ def unload_complete_models(*args):
         print(f'Unloaded {m.__class__.__name__} as complete.')
 
     gpu_complete_modules.clear()
-    torch.cuda.empty_cache()
+    torch.cuda.empty_cache() if torch.cuda.is_available() else None
     return
 
 

requirements.txt

@@ -16,3 +16,4 @@ einops
 opencv-contrib-python
 safetensors
 huggingface_hub
+spaces