Update app.py

app.py

@@ -14,6 +14,14 @@ CUSTOM_LORA_DIR = "./Custom_LoRAs"
 os.makedirs(LOCAL_LORA_DIR, exist_ok=True)
 os.makedirs(CUSTOM_LORA_DIR, exist_ok=True)
 
+# ------------------ DEVICE SETUP (✅ supports CPU-only Spaces) ------------------ #
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+print(f"🚀 Running on device: {device}")
+if device.type == "cpu":
+    print("⚠️ WARNING: No CUDA detected. Running on CPU. Generation may be slow.")
+
+# ------------------ Load Base LoRA ------------------ #
 print("downloading OmniConsistency base LoRA …")
 omni_consistency_path = hf_hub_download(
     repo_id="showlab/OmniConsistency",
@@ -23,11 +31,13 @@ omni_consistency_path = hf_hub_download(
 
 print("loading base pipeline …")
 pipe = FluxPipeline.from_pretrained(
-    BASE_PATH, torch_dtype=torch.bfloat16
-).to("cuda")
+    BASE_PATH, torch_dtype=dtype
+).to(device)
+
 set_single_lora(pipe.transformer, omni_consistency_path,
                 lora_weights=[1], cond_size=512)
 
+# ------------------ Download LoRA Styles ------------------ #
 def download_all_loras():
     lora_names = [
         "3D_Chibi", "American_Cartoon", "Chinese_Ink", "Clay_Toy",
@@ -48,7 +58,8 @@ def clear_cache(transformer):
     for _, attn_processor in transformer.attn_processors.items():
         attn_processor.bank_kv.clear()
 
-@spaces.GPU()
+# ------------------ Generation Function ------------------ #
+@spaces.GPU()  # Will fallback silently if GPU not available
 def generate_image(
     lora_name,        
     custom_repo_id,   
@@ -113,7 +124,7 @@ def generate_image(
     clear_cache(pipe.transformer)
     return uploaded_image, out_img
 
-# =============== Gradio UI ===============
+# ------------------ UI Interface ------------------ #
 def create_interface():
     demo_lora_names = [
         "3D_Chibi", "American_Cartoon", "Chinese_Ink", "Clay_Toy",
@@ -130,7 +141,6 @@ def create_interface():
       new_trigger = " ".join(lora_name.split("_"))+ " style,"
       return new_trigger + prompt
 
-    # Example data
     examples = [
         ["3D_Chibi", "", "3D Chibi style, Two smiling colleagues enthusiastically high-five in front of a whiteboard filled with technical notes about multimodal learning, reflecting a moment of success and collaboration at OpenAI.", 
         Image.open("./test_imgs/00.png"), 680, 1024, 3.5, 24, 42],