Update app.py

app.py

@@ -72,9 +72,10 @@ else:
     print("CUDA installation not found")
 
 
+device = torch.device('cuda')
 
 base_model = "black-forest-labs/FLUX.1-dev"
-pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16, token=huggingface_token)
+pipe = FluxPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16, token=huggingface_token).to(device)
 
 print('Loading and fusing lora, please wait...')
 pipe.load_lora_weights(hf_hub_download("gokaygokay/Flux-Game-Assets-LoRA-v2", "game_asst.safetensors"))
@@ -99,8 +100,6 @@ infer_config = config.infer_config
 
 IS_FLEXICUBES = True if config_name.startswith('instant-mesh') else False
 
-device = torch.device('cuda')
-
 # Load diffusion model for 3D generation
 print('Loading diffusion model ...')
 pipeline = DiffusionPipeline.from_pretrained(