dynamic gpu

app.py

@@ -13,12 +13,15 @@ pipe = None
 
 CHECKPOINT = "primecai/dsd_model"
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+
 def init_pipeline():
     global pipe
     transformer = FluxTransformer2DConditionalModel.from_pretrained(
         CHECKPOINT,
         subfolder="transformer",
-        torch_dtype=torch.bfloat16,
+        torch_dtype=dtype,
         low_cpu_mem_usage=False,
         ignore_mismatched_sizes=True,
         use_auth_token=os.getenv("HF_TOKEN"),
@@ -26,7 +29,7 @@ def init_pipeline():
     pipe = FluxConditionalPipeline.from_pretrained(
         "black-forest-labs/FLUX.1-dev",
         transformer=transformer,
-        torch_dtype=torch.bfloat16,
+        torch_dtype=dtype,
         use_auth_token=os.getenv("HF_TOKEN"),
     )
     pipe.load_lora_weights(
@@ -34,7 +37,7 @@ def init_pipeline():
         weight_name="pytorch_lora_weights.safetensors",
         use_auth_token=os.getenv("HF_TOKEN"),
     )
-    pipe.to("cuda")
+    pipe.to(device, dtype=dtype)
 
 
 def process_image_and_text(image, text, gemini_prompt, guidance, i_guidance, t_guidance):