Update app.py

app.py

@@ -111,7 +111,7 @@ def initialize_models(args):
     # Load main model and task head
     model = UnivaQwen2p5VLForConditionalGeneration.from_pretrained(
         args.model_path,
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
         attn_implementation="sdpa",
         quantization_config=quantization_config if args.nf4 else None,
     )
@@ -134,19 +134,19 @@ def initialize_models(args):
             args.flux_path,
             subfolder="text_encoder_2",
             quantization_config=quantization_config,
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float16,
         )
         pipe = FluxPipeline.from_pretrained(
             args.flux_path,
             transformer=model.denoise_tower.denoiser,
             text_encoder_2=text_encoder_2,
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float16,
         )
     else:
         pipe = FluxPipeline.from_pretrained(
             args.flux_path,
             transformer=model.denoise_tower.denoiser,
-            torch_dtype=torch.float32,
+            torch_dtype=torch.float16,
         )
     if args.offload:
         pipe.enable_model_cpu_offload()
@@ -159,7 +159,7 @@ def initialize_models(args):
     siglip_processor = SiglipImageProcessor.from_pretrained(args.siglip_path)
     siglip_model = SiglipVisionModel.from_pretrained(
         args.siglip_path,
-        torch_dtype=torch.float32,
+        torch_dtype=torch.float16,
     )
 
     return {