Update app.py

app.py

@@ -2,28 +2,30 @@
 
 from __future__ import annotations
 
+import requests
 import os
 import random
 
 import gradio as gr
 import numpy as np
-import PIL.Image
 import spaces
 import torch
-from diffusers import AutoencoderKL, DiffusionPipeline
+from PIL import Image
+from io import BytesIO
+from diffusers import AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image
 
-DESCRIPTION = "# SDXL"
+DESCRIPTION = "# Run any LoRA or SD Model"
 if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
+    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_REFINER = os.getenv("ENABLE_REFINER", "1") == "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_VAE", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -46,26 +48,37 @@ def generate(
     width: int = 1024,
     height: int = 1024,
     guidance_scale_base: float = 5.0,
-    guidance_scale_refiner: float = 5.0,
     num_inference_steps_base: int = 25,
-    num_inference_steps_refiner: int = 25,
+    strength_img2img: float = 0.7,
     use_vae: bool = False,
     use_lora: bool = False,
-    apply_refiner: bool = False,
-    model = 'cagliostrolab/animagine-xl-3.0',
-    vaecall = 'stabilityai/sd-vae-ft-mse',
-    lora = 'amazonaws-la/juliette',
+    model = 'stabilityai/stable-diffusion-xl-base-1.0',
+    vaecall = 'madebyollin/sdxl-vae-fp16-fix',
+    lora = '',
     lora_scale: float = 0.7,
-) -> PIL.Image.Image:
+    use_img2img: bool = False,
+    url = '',
+):
     if torch.cuda.is_available():
-
-        if not use_vae:  
-            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
             
-        if use_vae:
-            vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-            pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+        if not use_img2img:
+            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
+
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
             
+        if use_img2img:
+            pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
+
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+                    
+            response = requests.get(url)
+            init_image = Image.open(BytesIO(response.content)).convert("RGB")
+            init_image = init_image.resize((width, height))
+
         if use_lora:
             pipe.load_lora_weights(lora)
             pipe.fuse_lora(lora_scale)
@@ -88,7 +101,7 @@ def generate(
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
-    if not apply_refiner:
+    if not use_img2img:
         return pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -102,8 +115,10 @@ def generate(
             output_type="pil",
         ).images[0]
     else:
-        latents = pipe(
+        images = pipe(
             prompt=prompt,
+            image=init_image,
+            strength=strength_img2img,
             negative_prompt=negative_prompt,
             prompt_2=prompt_2,
             negative_prompt_2=negative_prompt_2,
@@ -112,75 +127,67 @@ def generate(
             guidance_scale=guidance_scale_base,
             num_inference_steps=num_inference_steps_base,
             generator=generator,
-            output_type="latent",
-        ).images
-        image = refiner(
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
-            guidance_scale=guidance_scale_refiner,
-            num_inference_steps=num_inference_steps_refiner,
-            image=latents,
-            generator=generator,
+            output_type="pil",
         ).images[0]
-        return image
-
+        return images
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
 ]
 
-with gr.Blocks(css="style.css") as demo:
-    gr.Markdown(DESCRIPTION)
-    gr.DuplicateButton(
-        value="Duplicate Space for private use",
-        elem_id="duplicate-button",
-        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
+    gr.HTML(
+        "<p><center>📙 For any additional support, join our <a href='https://discord.gg/JprjXpjt9K'>Discord</a></center></p>"
     )
+    gr.Markdown(DESCRIPTION, elem_id="description")
     with gr.Group():
-        model = gr.Text(label='Model')
-        vaecall = gr.Text(label='VAE')
-        lora = gr.Text(label='LoRA')
+        model = gr.Text(label='Model', placeholder='e.g. stabilityai/stable-diffusion-xl-base-1.0')
+        vaecall = gr.Text(label='VAE', placeholder='e.g. madebyollin/sdxl-vae-fp16-fix')
+        lora = gr.Text(label='LoRA', placeholder='e.g. nerijs/pixel-art-xl')
         lora_scale = gr.Slider(
+                info="The closer to 1, the more it will resemble LoRA, but errors may be visible.",
                 label="Lora Scale",
                 minimum=0.01,
                 maximum=1,
                 step=0.01,
                 value=0.7,
             )
+        url = gr.Text(label='URL (Img2Img)', placeholder='e.g https://example.com/image.png')
         with gr.Row():
             prompt = gr.Text(
+                placeholder="Input prompt",
                 label="Prompt",
                 show_label=False,
                 max_lines=1,
-                placeholder="Enter your prompt",
                 container=False,
             )
             run_button = gr.Button("Run", scale=0)
         result = gr.Image(label="Result", show_label=False)
     with gr.Accordion("Advanced options", open=False):
         with gr.Row():
+            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)
             use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
             use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
             use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False)
         negative_prompt = gr.Text(
+            placeholder="Input Negative Prompt",
             label="Negative prompt",
             max_lines=1,
-            placeholder="Enter a negative prompt",
             visible=False,
         )
         prompt_2 = gr.Text(
+            placeholder="Input Prompt 2",
             label="Prompt 2",
             max_lines=1,
-            placeholder="Enter your prompt",
             visible=False,
         )
         negative_prompt_2 = gr.Text(
+            placeholder="Input Negative Prompt 2",
             label="Negative prompt 2",
             max_lines=1,
-            placeholder="Enter a negative prompt",
             visible=False,
         )
 
@@ -207,38 +214,33 @@ with gr.Blocks(css="style.css") as demo:
                 step=32,
                 value=1024,
             )
-        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)
-        apply_refiner = gr.Checkbox(label="Apply refiner", value=False, visible=ENABLE_REFINER)
+            
         with gr.Row():
             guidance_scale_base = gr.Slider(
-                label="Guidance scale for base",
+                info="Scale for classifier-free guidance",
+                label="Guidance scale",
                 minimum=1,
                 maximum=20,
                 step=0.1,
                 value=5.0,
             )
+        with gr.Row():
             num_inference_steps_base = gr.Slider(
-                label="Number of inference steps for base",
+                info="Number of denoising steps",
+                label="Number of inference steps",
                 minimum=10,
                 maximum=100,
                 step=1,
                 value=25,
             )
-        with gr.Row(visible=False) as refiner_params:
-            guidance_scale_refiner = gr.Slider(
-                label="Guidance scale for refiner",
-                minimum=1,
-                maximum=20,
-                step=0.1,
-                value=5.0,
-            )
-            num_inference_steps_refiner = gr.Slider(
-                label="Number of inference steps for refiner",
-                minimum=10,
-                maximum=100,
-                step=1,
-                value=25,
+        with gr.Row():
+            strength_img2img = gr.Slider(
+                info="Strength for Img2Img",
+                label="Strength",
+                minimum=0,
+                maximum=1,
+                step=0.01,
+                value=0.7,
             )
 
     gr.Examples(
@@ -284,10 +286,10 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     )
-    apply_refiner.change(
+    use_img2img.change(
         fn=lambda x: gr.update(visible=x),
-        inputs=apply_refiner,
-        outputs=refiner_params,
+        inputs=use_img2img,
+        outputs=url,
         queue=False,
         api_name=False,
     )
@@ -319,16 +321,16 @@ with gr.Blocks(css="style.css") as demo:
             width,
             height,
             guidance_scale_base,
-            guidance_scale_refiner,
             num_inference_steps_base,
-            num_inference_steps_refiner,
+            strength_img2img,
             use_vae,
             use_lora,
-            apply_refiner,
             model,
             vaecall,
             lora,
             lora_scale,
+            use_img2img,
+            url,
         ],
         outputs=result,
         api_name="run",