Update app.py

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import spaces  # Import spaces for ZeroGPU compatibility
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 from janus.utils.io import load_pil_images
 import numpy as np
@@ -6,10 +7,6 @@ from PIL import Image
 from transformers import AutoConfig, AutoModelForCausalLM
 import torch
 
-##
-# Code from deepseek-ai/Janus
-# Space from huggingface/twodgirl.
-
 def generate(input_ids,
              width,
              height,
@@ -58,10 +55,10 @@ def unpack(dec, width, height, parallel_size=1):
     return visual_img
 
 @torch.inference_mode()
+@spaces.GPU  # Decorate the function for ZeroGPU compatibility
 def generate_image(prompt,
                    width,
                    height,
-                   # num_steps,
                    guidance,
                    seed):
     if seed > -1:
@@ -117,7 +114,6 @@ if __name__ == '__main__':
     model_path = 'deepseek-ai/Janus-1.3B'
     processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
     tokenizer = processor.tokenizer
-    # model: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
     config = AutoConfig.from_pretrained(model_path)
     language_config = config.language_config
     language_config._attn_implementation = 'eager'