update code

app.py

@@ -20,6 +20,7 @@ def extract_gen_content(text):
 
 def _load_model_processor():
 
+    device = "cuda" if torch.cuda.is_available() else "cpu"
     local_dir = snapshot_download(
         repo_id="xiabs/DreamOmni2",
         revision="main",
@@ -32,13 +33,14 @@ def _load_model_processor():
     pipe = DreamOmni2Pipeline.from_pretrained(
         "black-forest-labs/FLUX.1-Kontext-dev",
         torch_dtype=torch.bfloat16
-    )
+    ).to(device)
     pipe.load_lora_weights(lora_dir, adapter_name="edit")
     pipe.set_adapters(["edit"], adapter_weights=[1])    
 
     vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
         vlm_dir,
-        torch_dtype="bfloat16"
+        torch_dtype="bfloat16",
+        device=device
     )
     processor = AutoProcessor.from_pretrained(vlm_dir)
     return vlm_model, processor, pipe
@@ -59,7 +61,8 @@ def _launch_demo(vlm_model, processor, pipe):
         text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         image_inputs, video_inputs = process_vision_info(messages)
         inputs = processor(text=[text], images=image_inputs, videos=video_inputs, padding=True, return_tensors="pt")
-        inputs = inputs.to("cuda")
+        inputs = inputs.to(device=vlm_model.device)
+        print(vlm_model.device, '++++')
 
         generated_ids = vlm_model.generate(**inputs, do_sample=False, max_new_tokens=4096)
         generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
@@ -116,7 +119,7 @@ def _launch_demo(vlm_model, processor, pipe):
         image.save(output_path)
         print(f"Edit result saved to {output_path}")
 
-
+    @spaces.GPU()
     def process_request(image_file_1, image_file_2, instruction):
         # debugpy.listen(5678)
         # print("Waiting for debugger attach...")