space

app.py

@@ -1,4 +1,5 @@
 import argparse
+import spaces
 from visualcloze import VisualClozeModel
 import gradio as gr
 import demo_tasks
@@ -313,7 +314,8 @@ def create_demo(model):
                     images[i].append(inputs[i * max_grid_w + j])
             seed, cfg, steps, upsampling_steps, upsampling_noise, layout_text, task_text, content_text = inputs[-8:]
                 
-            results = model.process_images(
+            results = generate(
+                model, 
                 images, 
                 [layout_text, task_text, content_text], 
                 seed=seed, cfg=cfg, steps=steps, 
@@ -458,6 +460,11 @@ def create_demo(model):
         
     return demo
 
+
+@spaces.GPU
+def generate(model, **args):
+    return model.process_images(args)
+
 def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_path", type=str, default="models/visualcloze-384-lora.pth")

visualcloze.py

@@ -1,6 +1,5 @@
 
 import random
-import spaces
 from einops import rearrange
 from diffusers.models import AutoencoderKL
 from PIL import Image
@@ -75,7 +74,6 @@ 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,