Update app.py

app.py

@@ -13,8 +13,9 @@ import torch
 from diffusers import DiffusionPipeline
 from typing import Tuple
 
-# Check if GPU is available; fallback to CPU if needed
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+# Initialize device to None
+device = None
+pipe = None
 
 # Setup rules for bad words (ensure the prompts are kid-friendly)
 bad_words = json.loads(os.getenv('BAD_WORDS', '["violence", "blood", "scary", "death", "ghost"]'))
@@ -68,21 +69,21 @@ DESCRIPTION = """## Children's Sticker Generator
 Generate fun and playful stickers for children using AI.
 """
 
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>⚠️Running on CPU, This may not work on CPU.</p>"
-
 MAX_SEED = np.iinfo(np.int32).max
 CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES", "0") == "1"
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+def initialize_pipeline(device_type):
+    global device, pipe
+    device = torch.device(device_type)
+    pipe = DiffusionPipeline.from_pretrained(
+        "SG161222/RealVisXL_V3.0_Turbo",
+        torch_dtype=torch.float32 if device_type == "cpu" else torch.float16,
+        use_safetensors=True,
+        variant="fp32" if device_type == "cpu" else "fp16"
+    ).to(device)
 
-# Initialize the DiffusionPipeline
-pipe = DiffusionPipeline.from_pretrained(
-    "SG161222/RealVisXL_V3.0_Turbo",  # or any model of your choice
-    torch_dtype=torch.float16,
-    use_safetensors=True,
-    variant="fp16"
-).to(device)
+# Initialize with CPU by default
+initialize_pipeline("cpu")
 
 # Convert mm to pixels for a specific DPI (300) and ensure divisible by 8
 def mm_to_pixels(mm, dpi=300):
@@ -134,8 +135,15 @@ def generate(
     guidance_scale: float = 3,
     randomize_seed: bool = False,
     background: str = "transparent",
+    device_type: str = "cpu",
     progress=gr.Progress(track_tqdm=True),
 ):
+    global device, pipe
+    
+    # Switch device if necessary
+    if device.type != device_type:
+        initialize_pipeline(device_type)
+
     if check_text(prompt, negative_prompt):
         raise ValueError("Prompt contains restricted words.")
     
@@ -145,7 +153,7 @@ def generate(
     # Apply style
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
     # Ensure we have only white or transparent background options
     width, height = size_map.get(size, (1024, 1024))
@@ -241,6 +249,11 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             step=0.1,
             value=15.7,
         )
+        device_selection = gr.Radio(
+            choices=["cpu", "cuda"],
+            value="cpu",
+            label="Device",
+        )
 
     gr.Examples(
         examples=examples,
@@ -267,6 +280,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             guidance_scale,
             randomize_seed,
             background_selection,
+            device_selection,
         ],
         outputs=[result, seed],
         api_name="run",