Update app.py

app.py

@@ -6,11 +6,12 @@ import gradio as gr
 import numpy as np
 import torch
 from PIL import Image
-from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffusionXLInstructPix2PixPipeline, AutoencoderKL
 from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download
 from huggingface_hub import InferenceClient
 
+
 help_text = """
 To optimize image results:
 - Adjust the **Image CFG weight** if the image isn't changing enough or is changing too much. Lower it to allow bigger changes, or raise it to preserve original details.
@@ -40,7 +41,6 @@ normal_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl.safet
 
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
 
-# Use a pre-converted fp16 VAE for faster loading and inference
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
 pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
@@ -49,14 +49,21 @@ pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
 pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
 pipe_edit.to("cuda")
 
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
+
+if not torch.cuda.is_available():
+    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
 # Image Generator
-# Keep the models loaded globally for reuse
 if torch.cuda.is_available():
     pipe = StableDiffusionXLPipeline.from_pretrained(
         "fluently/Fluently-XL-v4",
         torch_dtype=torch.float16,
         use_safetensors=True,
-    ).to("cuda")
+    )
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
     pipe.set_adapters("dalle")
@@ -87,6 +94,7 @@ def king(type = "Image Generation",
         text_cfg_scale = text_cfg_scale
         image_cfg_scale = image_cfg_scale
         input_image = input_image
+
         steps=steps
         generator = torch.manual_seed(seed)
         output_image = pipe_edit(
@@ -95,9 +103,10 @@ def king(type = "Image Generation",
             num_inference_steps=steps, generator=generator).images[0]
         return seed, output_image
     else :
+        pipe.to(device)
         seed = int(randomize_seed_fn(seed, randomize_seed))
-        generator = torch.Generator().manual_seed(seed) # Move generator to cuda for speed
-
+        generator = torch.Generator().manual_seed(seed)   
+        
         options = {
             "prompt":instruction,
             "width":width,
@@ -116,6 +125,7 @@ def king(type = "Image Generation",
 def response(instruction, input_image=None):
     if input_image is None:
         output="Image Generation"
+        yield output
     else:
         client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")