Update app.py

app.py

@@ -47,14 +47,24 @@ if torch.cuda.is_available():
         add_watermarker=False,
         variant="fp16"
     )
+    pipe2 = DiffusionPipeline.from_pretrained(
+        "SG161222/RealVisXL_V4.0",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+        add_watermarker=False,
+        variant="fp16"
+    )
     if ENABLE_CPU_OFFLOAD:
         pipe.enable_model_cpu_offload()
+        pipe2.enable_model_cpu_offload()
     else:
         pipe.to(device)    
+        pipe2.to(device)    
         print("Loaded on Device!")
     
     if USE_TORCH_COMPILE:
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+        pipe2.unet = torch.compile(pipe2.unet, mode="reduce-overhead", fullgraph=True)
         print("Model Compiled!")
 
 
@@ -70,33 +80,9 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     return seed
 
 
-def generate(prompt: str,
-    negative_prompt: str = "",
-    use_negative_prompt: bool = False,
-    seed: int = 0,
-    width: int = 1024,
-    height: int = 1024,
-    guidance_scale: float = 3,
-    randomize_seed: bool = False,
-    use_resolution_binning: bool = True,
-    progress=gr.Progress(track_tqdm=True)):
-    if check_text(prompt, negative_prompt):
-        return 'NSFW Detection'
-    return _generate(
-        prompt,
-        negative_prompt,
-        use_negative_prompt,
-        seed,
-        width,
-        height,
-        guidance_scale,
-        randomize_seed,
-        use_resolution_binning,
-        progress
-    )
 
 @spaces.GPU(enable_queue=True)
-def _generate(
+def generate(
     prompt: str,
     negative_prompt: str = "",
     use_negative_prompt: bool = False,
@@ -115,18 +101,22 @@ def _generate(
     if not use_negative_prompt:
         negative_prompt = ""  # type: ignore
     negative_prompt += default_negative    
-    images = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        width=width,
-        height=height,
-        guidance_scale=guidance_scale,
-        num_inference_steps=25,
-        generator=generator,
-        num_images_per_prompt=NUM_IMAGES_PER_PROMPT,
-        use_resolution_binning=use_resolution_binning,
-        output_type="pil",
-    ).images
+
+    options = {
+        "prompt":prompt,
+        "negative_prompt":negative_prompt,
+        "width":width,
+        "height":height,
+        "guidance_scale":guidance_scale,
+        "num_inference_steps":25,
+        "generator":generator,
+        "num_images_per_prompt":NUM_IMAGES_PER_PROMPT,
+        "use_resolution_binning":use_resolution_binning,
+        "output_type":"pil",
+
+    }
+    
+    images = pipe(**options).images+pipe2(**options).images
 
     image_paths = [save_image(img) for img in images]
     return image_paths, seed