Update app.py

app.py

@@ -24,6 +24,8 @@ subprocess.run(
     shell=True,
 )
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # Model paths
 EMU_GEN_HUB = "BAAI/Emu3-Gen"
 EMU_CHAT_HUB = "BAAI/Emu3-Chat"
@@ -33,28 +35,28 @@ VQ_HUB = "BAAI/Emu3-VisionTokenizer"
 # Emu3-Gen model and processor
 gen_model = AutoModelForCausalLM.from_pretrained(
     EMU_GEN_HUB,
-    device_map="cuda:0",
+    device_map="cpu",
     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(
     EMU_CHAT_HUB,
-    device_map="cuda:0",
+    device_map="cpu",
     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(
     VQ_HUB, trust_remote_code=True
 )
 image_tokenizer = AutoModel.from_pretrained(
-    VQ_HUB, device_map="cuda:0", trust_remote_code=True
-).eval()
+    VQ_HUB, device_map="cpu", trust_remote_code=True
+).eval().to(device)
 processor = Emu3Processor(
     image_processor, image_tokenizer, tokenizer
 )
@@ -97,7 +99,7 @@ def generate_image(prompt):
             UnbatchedClassifierFreeGuidanceLogitsProcessor(
                 classifier_free_guidance,
                 gen_model,
-                unconditional_ids=neg_inputs.input_ids.to("cuda:0"),
+                unconditional_ids=neg_inputs.input_ids.to(device),
             ),
             PrefixConstrainedLogitsProcessor(
                 constrained_fn,
@@ -108,7 +110,7 @@ def generate_image(prompt):
 
     # Generate
     outputs = gen_model.generate(
-        pos_inputs.input_ids.to("cuda:0"),
+        pos_inputs.input_ids.to(device),
         generation_config=GENERATION_CONFIG,
         logits_processor=logits_processor,
     )
@@ -139,7 +141,7 @@ def vision_language_understanding(image, text):
 
     # Generate
     outputs = chat_model.generate(
-        inputs.input_ids.to("cuda:0"),
+        inputs.input_ids.to(device),
         generation_config=GENERATION_CONFIG,
         max_new_tokens=320,
     )