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