Update app.py

app.py

@@ -24,6 +24,8 @@ subprocess.run(
     shell=True,
 )
 
+print(gr.__version__)
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Model paths
@@ -39,7 +41,7 @@ gen_model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16,
     attn_implementation="flash_attention_2",
     trust_remote_code=True,
-).to(device)
+)
 
 # Emu3-Chat model and processor
 chat_model = AutoModelForCausalLM.from_pretrained(
@@ -48,7 +50,7 @@ chat_model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16,
     attn_implementation="flash_attention_2",
     trust_remote_code=True,
-).to(device)
+)
 
 tokenizer = AutoTokenizer.from_pretrained(EMU_CHAT_HUB, trust_remote_code=True)
 image_processor = AutoImageProcessor.from_pretrained(
@@ -56,11 +58,16 @@ image_processor = AutoImageProcessor.from_pretrained(
 )
 image_tokenizer = AutoModel.from_pretrained(
     VQ_HUB, device_map="cpu", trust_remote_code=True
-).eval().to(device)
+).eval()
 processor = Emu3Processor(
     image_processor, image_tokenizer, tokenizer
 )
 
+print(device)
+gen_model.to(device)
+chat_model.to(device)
+image_tokenizer.to(device)
+
 @spaces.GPU(duration=300)
 def generate_image(prompt):
     POSITIVE_PROMPT = " masterpiece, film grained, best quality."