Update app.py

app.py

@@ -98,15 +98,12 @@ def load_and_prepare_model(model_id):
         torch_dtype=torch.bfloat16,
         use_safetensors=True,
         add_watermarker=False,
-    ).to(device)
+    )
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     
     if USE_TORCH_COMPILE:
         pipe.compile()
     
-    if ENABLE_CPU_OFFLOAD:
-        pipe.enable_model_cpu_offload()
-    
     return pipe
 
 # Preload and compile both models
@@ -124,7 +121,7 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-@spaces.GPU(duration=40, enable_queue=True)
+@spaces.GPU(duration=70, enable_queue=True)
 def generate(
     model_choice: str,
     prompt: str,
@@ -143,7 +140,57 @@ def generate(
 ):
     global models
     pipe = models[model_choice]
-    
+    pipe.to("cuda:0")
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    prompt, negative_prompt = apply_style(style_selection, prompt, negative_prompt)
+
+    options = {
+        "prompt": [prompt] * num_images,
+        "negative_prompt": [negative_prompt] * num_images if use_negative_prompt else None,
+        "width": width,
+        "height": height,
+        "guidance_scale": guidance_scale,
+        "num_inference_steps": num_inference_steps,
+        "generator": generator,
+        "output_type": "pil",
+    }
+
+    if use_resolution_binning:
+        options["use_resolution_binning"] = True
+
+    images = []
+    for i in range(0, num_images, BATCH_SIZE):
+        batch_options = options.copy()
+        batch_options["prompt"] = options["prompt"][i:i+BATCH_SIZE]
+        if "negative_prompt" in batch_options:
+            batch_options["negative_prompt"] = options["negative_prompt"][i:i+BATCH_SIZE]
+        images.extend(pipe(**batch_options).images)
+
+    image_paths = [save_image(img) for img in images]
+    return image_paths, seed
+
+def generate_cpu(
+    model_choice: str,
+    prompt: str,
+    negative_prompt: str = "",
+    use_negative_prompt: bool = False,
+    style_selection: str = DEFAULT_STYLE_NAME,
+    seed: int = 1,
+    width: int = 768,
+    height: int = 768,
+    guidance_scale: float = 4,
+    num_inference_steps: int = 150,
+    randomize_seed: bool = False,
+    use_resolution_binning: bool = True, 
+    num_images: int = 1,  
+    progress=gr.Progress(track_tqdm=True),
+):
+    global models
+    pipe = models[model_choice]
+    pipe.to("cpu")
+
     seed = int(randomize_seed_fn(seed, randomize_seed))
     generator = torch.Generator(device=device).manual_seed(seed)
 
@@ -215,6 +262,7 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
             container=False,
         )
         run_button = gr.Button("Run", scale=0)
+        cpu_run_button = gr.Button("CPU Run", scale=0)
     result = gr.Gallery(label="Result", columns=1, show_label=False) 
 
     with gr.Row():
@@ -328,6 +376,29 @@ with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
         outputs=[result, seed],
     )
 
+    gr.on(
+        triggers=[
+            cpu_run_button.click,
+        ],
+        api_name="generate",  # Add this line
+        fn=generate_cpu,
+        inputs=[
+            model_choice,
+            prompt,
+            negative_prompt,
+            use_negative_prompt,
+            style_selection,
+            seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+            randomize_seed, 
+            num_images, 
+        ],
+        outputs=[result, seed],
+    )
+
     gr.Markdown("### REALVISXL V5.0")
     predefined_gallery = gr.Gallery(label="REALVISXL V5.0", columns=3, show_label=False, value=load_predefined_images1())