Update app.py

app.py

@@ -12,21 +12,20 @@ import spaces
 import torch
 from PIL import Image
 from io import BytesIO
-from diffusers.utils import load_image
-from diffusers import AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting
+from diffusers import AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image
 
 DESCRIPTION = "# Run any LoRA or SD Model"
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>⚠️ This space is running on the CPU. This demo doesn't work on CPU 😞! Run on a GPU by duplicating this space or test our website for free and unlimited by <a href='https://squaadai.com'>clicking here</a>, which provides these and more options.</p>"
 
 MAX_SEED = np.iinfo(np.int32).max
+CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1824"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
 ENABLE_USE_VAE = os.getenv("ENABLE_USE_VAE", "1") == "1"
-ENABLE_USE_IMG2IMG = os.getenv("ENABLE_USE_IMG2IMG", "1") == "1"
-ENABLE_USE_INPAINTING = os.getenv("ENABLE_USE_INPAINTING", "1") == "1"
+ENABLE_USE_IMG2IMG = os.getenv("ENABLE_USE_VAE", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -58,10 +57,7 @@ def generate(
     lora = '',
     lora_scale: float = 0.7,
     use_img2img: bool = False,
-    use_inpainting: bool = False,
     url = '',
-    img_url = '',
-    mask_url = '',
 ):
     if torch.cuda.is_available():
             
@@ -78,17 +74,11 @@ def generate(
             if use_vae:  
                 vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
                 pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
-                
-        if use_inpainting:
-            pipe = AutoPipelineForInpainting.from_pretrained(model, torch_dtype=torch.float16)
-            
+                    
             response = requests.get(url)
             init_image = Image.open(BytesIO(response.content)).convert("RGB")
             init_image = init_image.resize((width, height))
 
-            image_init = load_image(img_url)
-            mask_image = load_image(mask_url)
-            
         if use_lora:
             pipe.load_lora_weights(lora)
             pipe.fuse_lora(lora_scale)
@@ -111,23 +101,20 @@ def generate(
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
-    if use_inpainting:
-        image = pipe(
+    if not use_img2img:
+        return pipe(
             prompt=prompt,
-            image=image_init,
-            mask_image=mask_image,
-            strength=strength_img2img,
             negative_prompt=negative_prompt,
             prompt_2=prompt_2,
+            negative_prompt_2=negative_prompt_2,
             width=width,
             height=height,
-            negative_prompt_2=negative_prompt_2,
             guidance_scale=guidance_scale_base,
             num_inference_steps=num_inference_steps_base,
             generator=generator,
+            output_type="pil",
         ).images[0]
-        return image
-    elif use_img2img:
+    else:
         images = pipe(
             prompt=prompt,
             image=init_image,
@@ -143,19 +130,11 @@ def generate(
             output_type="pil",
         ).images[0]
         return images
-    else:
-        return pipe(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
-            width=width,
-            height=height,
-            guidance_scale=guidance_scale_base,
-            num_inference_steps=num_inference_steps_base,
-            generator=generator,
-            output_type="pil",
-        ).images[0]
+
+examples = [
+    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    "An astronaut riding a green horse",
+]
 
 with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     gr.HTML(
@@ -175,8 +154,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
                 value=0.7,
             )
         url = gr.Text(label='URL (Img2Img)', placeholder='e.g https://example.com/image.png')
-        img_url = gr.Text(label='URL (Image Inpainting)', placeholder='e.g https://example.com/image.png')
-        mask_url = gr.Text(label='URL (Mask Image Inpainting)', placeholder='e.g https://example.com/image.png')
         with gr.Row():
             prompt = gr.Text(
                 placeholder="Input prompt",
@@ -189,7 +166,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         result = gr.Image(label="Result", show_label=False)
     with gr.Accordion("Advanced options", open=False):
         with gr.Row():
-            use_inpainting = gr.Checkbox(label='Use Inpainting', value=False, visible=ENABLE_USE_INPAINTING)
             use_img2img = gr.Checkbox(label='Use Img2Img', value=False, visible=ENABLE_USE_IMG2IMG)
             use_vae = gr.Checkbox(label='Use VAE', value=False, visible=ENABLE_USE_VAE)
             use_lora = gr.Checkbox(label='Use Lora', value=False, visible=ENABLE_USE_LORA)
@@ -267,6 +243,14 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
                 value=0.7,
             )
 
+    gr.Examples(
+        examples=examples,
+        inputs=prompt,
+        outputs=result,
+        fn=generate,
+        cache_examples=CACHE_EXAMPLES,
+    )
+
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
@@ -309,20 +293,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         queue=False,
         api_name=False,
     )
-    use_inpainting.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_inpainting,
-        outputs=img_url,
-        queue=False,
-        api_name=False,
-    )
-    use_inpainting.change(
-        fn=lambda x: gr.update(visible=x),
-        inputs=use_inpainting,
-        outputs=mask_url,
-        queue=False,
-        api_name=False,
-    )
 
     gr.on(
         triggers=[
@@ -360,10 +330,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
             lora,
             lora_scale,
             use_img2img,
-            use_inpainting,
             url,
-            img_url,
-            mask_url,
         ],
         outputs=result,
         api_name="run",