Fix Boogu FP8 loading on ZeroGPU

app.py

@@ -2,6 +2,7 @@ import gc
 import math
 import os
 import random
+import shutil
 import time
 from threading import Lock
 
@@ -12,6 +13,7 @@ os.environ.setdefault("device", "cuda:0")
 import gradio as gr
 import spaces
 import torch
+from huggingface_hub import snapshot_download
 from PIL import ImageOps
 
 from boogu.models.transformers.transformer_boogu import BooguImageTransformer2DModel
@@ -91,12 +93,28 @@ def cleanup_pipe():
 
 
 def load_fp8_transformer(repo_id, dtype):
+    repo_dir = prepare_repo_dir(repo_id)
     return BooguImageTransformer2DModel.from_pretrained(
-        repo_id,
-        subfolder="transformer",
+        os.path.join(repo_dir, "transformer"),
         torch_dtype=dtype,
         use_safetensors=False,
+    ), repo_dir
+
+
+def prepare_repo_dir(repo_id):
+    repo_dir = snapshot_download(
+        repo_id,
+        token=os.environ.get("HF_TOKEN") or None,
+    )
+    transformer_dir = os.path.join(repo_dir, "transformer")
+    source = os.path.join(transformer_dir, "transformer_boogu.py")
+    compat = os.path.join(
+        transformer_dir,
+        "boogu.models.transformers.transformer_boogu.py",
     )
+    if os.path.exists(source) and not os.path.exists(compat):
+        shutil.copyfile(source, compat)
+    return repo_dir
 
 
 def get_pipe(mode):
@@ -121,9 +139,9 @@ def get_pipe(mode):
 
         print(f"Loading {repo_id}...", flush=True)
         started = time.perf_counter()
-        transformer = load_fp8_transformer(repo_id, dtype)
+        transformer, repo_dir = load_fp8_transformer(repo_id, dtype)
         pipe = pipeline_cls.from_pretrained(
-            repo_id,
+            repo_dir,
             torch_dtype=dtype,
             trust_remote_code=True,
             transformer=transformer,

requirements.txt

@@ -1,4 +1,3 @@
 gradio[mcp]==5.49.1
 spaces
-diffusers[torch]==0.35.2
 git+https://github.com/boogu-project/Boogu-Image.git@8e329714727e6fb0804731b81681a121d1698f91