:tada: change process

app.py

@@ -17,23 +17,14 @@ DIFFUSERS_MODEL_IDS = [
     # Other Models
     "Prgckwb/trpfrog-diffusion",
 ]
-
 EXTERNAL_MODEL_MAPPING = {
     "Beautiful Realistic Asians": "checkpoints/diffusers/Beautiful Realistic Asians v7",
 }
-
 MODEL_CHOICES = DIFFUSERS_MODEL_IDS + list(EXTERNAL_MODEL_MAPPING.keys())
 
 current_model_id = "stabilityai/stable-diffusion-3-medium-diffusers"
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
-if device == 'cuda':
-    dtype = torch.float16
-    pipe = DiffusionPipeline.from_pretrained(
-        current_model_id,
-        torch_dtype=dtype,
-    )
-    pipe.enable_model_cpu_offload()
+pipe = None
 
 
 @dataclasses.dataclass
@@ -86,6 +77,7 @@ def inference(
         num_inference_steps: int = 50,
         num_images: int = 4,
         safety_checker: bool = True,
+        use_model_offload:bool = False,
         progress=gr.Progress(track_tqdm=True),
 ) -> Image.Image:
     progress(0, "Starting inference...")
@@ -100,9 +92,8 @@ def inference(
 
         pipe = DiffusionPipeline.from_pretrained(
             model_id,
-            torch_dtype=dtype,
+            torch_dtype=torch.float16,
         )
-        pipe.enable_model_cpu_offload()
 
         current_model_id = model_id
 
@@ -118,6 +109,11 @@ def inference(
 
     # Generation
     progress(0.4, 'Generating images...')
+    if use_model_offload:
+        pipe.enable_model_cpu_offload()
+    else:
+        pipe = pipe.to('cuda')
+
     images = pipe(
         prompt,
         negative_prompt=negative_prompt,
@@ -168,6 +164,7 @@ if __name__ == "__main__":
 
                     with gr.Row():
                         safety_checker = gr.Checkbox(value=True, label='Use Safety Checker')
+                        model_offload = gr.Checkbox(value=False, label='Use Model Offload')
 
             with gr.Column():
                 output_image = gr.Image(label="Image", type="pil")
@@ -181,7 +178,8 @@ if __name__ == "__main__":
             guidance_scale,
             num_inference_step,
             num_images,
-            safety_checker
+            safety_checker,
+            model_offload,
         ]
 
         btn = gr.Button("Generate")