space

app.py

@@ -1,10 +1,9 @@
 import argparse
-import spaces
+from visualcloze import VisualClozeModel
 import gradio as gr
 import demo_tasks
 from functools import partial
 from data.prefix_instruction import get_layout_instruction
-from visualcloze import VisualClozeModel
 from huggingface_hub import snapshot_download
 
 
@@ -306,7 +305,6 @@ def create_demo(model):
             updates.append(gr.update(elem_id="layout_prompt", value=get_layout_instruction(w, actual_h)))
             return updates
 
-        @spaces.GPU
         def generate_image(*inputs):
             images = []
             for i in range(model.grid_h):
@@ -467,6 +465,7 @@ def parse_args():
     parser.add_argument("--resolution", type=int, default=384)
     return parser.parse_args()
 
+
 if __name__ == "__main__":
     args = parse_args()
     

visualcloze.py

@@ -1,5 +1,6 @@
 
 import random
+import spaces
 from einops import rearrange
 from diffusers.models import AutoencoderKL
 from PIL import Image
@@ -12,7 +13,6 @@ from flux.util import load_clip, load_t5, load_flow_model
 from transport import Sampler, create_transport
 from imgproc import to_rgb_if_rgba
 
-
 def center_crop(image, target_size):
     width, height = image.size
     new_width, new_height = target_size
@@ -75,6 +75,7 @@ def resize_with_aspect_ratio(img, resolution, divisible=16, aspect_ratio=None):
         return img.resize((new_w, new_h), Image.LANCZOS)
 
 
+@spaces.GPU
 class VisualClozeModel:
     def __init__(
         self, model_path, model_name="flux-dev-fill-lora", max_length=512, lora_rank=256,