Update app.py

app.py

@@ -98,12 +98,10 @@ def img2b64(image_path):
     data_uri = f"data:image/jpeg;base64,{b64}"
     return data_uri
 
-@spaces.GPU(duration=900)
 @spaces.GPU
 def initialize_models(args):
     os.makedirs("tmp", exist_ok=True)
     # Paths
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     quantization_config = BitsAndBytesConfig(
         load_in_4bit=True,
@@ -114,16 +112,16 @@ def initialize_models(args):
     # Load main model and task head
     model = UnivaQwen2p5VLForConditionalGeneration.from_pretrained(
         args.model_path,
-        torch_dtype=torch.bfloat16,
+        torch_dtype=torch.float32,
         attn_implementation="sdpa",
         quantization_config=quantization_config if args.nf4 else None,
-    ).to(device)
+    )
     task_head = nn.Sequential(
         nn.Linear(3584, 10240),
         nn.SiLU(),
         nn.Dropout(0.3),
         nn.Linear(10240, 2)
-    ).to(device)
+    )
     task_head.load_state_dict(torch.load(os.path.join(args.model_path, 'task_head_final.pt')))
     task_head.eval()
 
@@ -137,20 +135,20 @@ def initialize_models(args):
             args.flux_path,
             subfolder="text_encoder_2",
             quantization_config=quantization_config,
-            torch_dtype=torch.bfloat16,
+            torch_dtype=torch.float32,
         )
         pipe = FluxPipeline.from_pretrained(
             args.flux_path,
             transformer=model.denoise_tower.denoiser,
             text_encoder_2=text_encoder_2,
-            torch_dtype=torch.bfloat16,
-        ).to(device)
+            torch_dtype=torch.float32,
+        )
     else:
         pipe = FluxPipeline.from_pretrained(
             args.flux_path,
             transformer=model.denoise_tower.denoiser,
-            torch_dtype=torch.bfloat16,
-        ).to(device)
+            torch_dtype=torch.float32,
+        )
     if args.offload:
         pipe.enable_model_cpu_offload()
         pipe.enable_vae_slicing()
@@ -162,8 +160,8 @@ def initialize_models(args):
     siglip_processor = SiglipImageProcessor.from_pretrained(args.siglip_path)
     siglip_model = SiglipVisionModel.from_pretrained(
         args.siglip_path,
-        torch_dtype=torch.bfloat16,
-    ).to(device)
+        torch_dtype=torch.float32,
+    )
 
     return {
         'model': model,
@@ -174,12 +172,23 @@ def initialize_models(args):
         'text_encoders': text_encoders,
         'siglip_processor': siglip_processor,
         'siglip_model': siglip_model,
-        'device': device,
+        
     }
 
+@spaces.GPU(duration=600)
+def to_device(state):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    state['model'] = state['model'].to(device, dtype=torch.bfloat16)
+    state['task_head'] = state['task_head'].to(device, dtype=torch.bfloat16)
+    state['pipe'] = state['pipe'].to(device, dtype=torch.bfloat16)
+    state['text_encoders'] = state['text_encoders'].to(device, dtype=torch.bfloat16)
+    state['siglip_model'] = state['siglip_model'].to(device, dtype=torch.bfloat16)
+    state['device'] = device
+    return state
 
 args = parse_args()
 state = initialize_models(args)
+state = to_device(state)
 
 @spaces.GPU
 def process_large_image(raw_img):