Update app.py

app.py

@@ -20,8 +20,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 pipe_sd3 = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
 pipe_flux = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
-pipe_sd3.to(device)
-pipe_flux.to(device)
+# pipe_sd3.to(device)
+# pipe_flux.to(device)
 
 # scheduler = pipe.scheduler
 # pipe = pipe.to(device)
@@ -97,15 +97,15 @@ def FlowEditRun(
 
     if model_type == 'FLUX':
         # pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
-        pipe = pipe_flux 
+        pipe = pipe_flux.to(device) 
     elif model_type == 'SD3':
         # pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, token=os.getenv('HF_ACCESS_TOK'))
-        pipe = pipe_sd3
+        pipe = pipe_sd3.to(device) 
     else:
         raise NotImplementedError(f"Model type {model_type} not implemented")
 
         scheduler = pipe.scheduler
-        pipe = pipe.to(device)
+        # pipe = pipe.to(device)