ZeroGPU fixes

app.py

@@ -1,16 +1,16 @@
-import torch
 import spaces
+import torch
 from model_loader import load_model_and_processor
 from image_generator import process_and_generate
 from gradio_interface import create_gradio_interface
 
 if __name__ == "__main__":
-    # Set the model path
-    model_path = "deepseek-ai/Janus-1.3B"
+    import subprocess
+    subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
+
+
 
-    # Load the model and processor
-    vl_gpt, vl_chat_processor = load_model_and_processor(model_path)
 
     # Create and launch the Gradio interface
-    demo = create_gradio_interface(vl_gpt, vl_chat_processor, process_and_generate)
+    demo = create_gradio_interface(process_and_generate)
     demo.launch(allowed_paths=["/"])

gradio_interface.py

@@ -1,9 +1,10 @@
 import gradio as gr
 from PIL import Image
 
-def create_gradio_interface(vl_gpt, vl_chat_processor, process_and_generate):
+
+def create_gradio_interface(process_and_generate):
     def gradio_process_and_generate(input_image, prompt, num_images, cfg_weight):
-        return process_and_generate(vl_gpt, vl_chat_processor, input_image, prompt, num_images, cfg_weight)
+        return process_and_generate(input_image, prompt, num_images, cfg_weight)
 
     explanation = """Janus 1.3B uses a differerent visual encoder for understanding and generation.
 

image_generator.py

@@ -3,9 +3,10 @@ import PIL.Image
 import torch
 import numpy as np
 from janus.utils.io import load_pil_images
+from model_loader import load_model_and_processor
 from janus.models import MultiModalityCausalLM, VLChatProcessor
 from functools import lru_cache
-
+import spaces
 
 def prepare_classifier_free_guidance_input(input_embeds, vl_chat_processor, mmgpt, batch_size=16):
     uncond_input_ids = torch.full((1, input_embeds.shape[1]), 
@@ -26,7 +27,6 @@ def prepare_classifier_free_guidance_input(input_embeds, vl_chat_processor, mmgp
     
     return combined_input_embeds
 
-@spaces.GPU
 @torch.inference_mode()
 def generate(
     mmgpt: MultiModalityCausalLM,
@@ -83,7 +83,14 @@ def get_start_tag_embed(vl_gpt, vl_chat_processor):
             vl_chat_processor.tokenizer.encode(vl_chat_processor.image_start_tag, add_special_tokens=False, return_tensors="pt").to(vl_gpt.device)
         )
     
-def process_and_generate(vl_gpt, vl_chat_processor, input_image, prompt, num_images=4, cfg_weight=5):
+@spaces.GPU
+def process_and_generate(input_image, prompt, num_images=4, cfg_weight=5):
+    # Set the model path
+    model_path = "deepseek-ai/Janus-1.3B"
+
+    # Load the model and processor
+    vl_gpt, vl_chat_processor = load_model_and_processor(model_path)
+    
     start_tag_embed = get_start_tag_embed(vl_gpt, vl_chat_processor)
 
     nl = '\n'

model_loader.py

@@ -1,7 +1,9 @@
 import torch
 from transformers import AutoModelForCausalLM
 from janus.models import MultiModalityCausalLM, VLChatProcessor
+import spaces
 
+@spaces.GPU
 def load_model_and_processor(model_path):
     vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_path)
     tokenizer = vl_chat_processor.tokenizer