Update app.py

app.py

@@ -2,17 +2,15 @@
 
 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 PIL import Image
-from io import BytesIO
-from diffusers import AutoencoderKL, DiffusionPipeline, StableDiffusionImg2ImgPipeline
+from diffusers import AutoencoderKL, DiffusionPipeline
 
 DESCRIPTION = "# Run any LoRA or SD Model"
 if not torch.cuda.is_available():
@@ -50,22 +48,16 @@ def generate(
     num_inference_steps_base: int = 25,
     use_vae: bool = False,
     use_lora: bool = False,
-    model = 'runwayml/stable-diffusion-v1-5',
+    model = 'stabilityai/stable-diffusion-xl-base-1.0',
     vaecall = 'madebyollin/sdxl-vae-fp16-fix',
     lora = '',
     lora_scale: float = 0.7,
-):
+) -> PIL.Image.Image:
     if torch.cuda.is_available():
 
         if not use_vae:  
-            pipe = StableDiffusionImg2ImgPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
+            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
             
-            url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"
-
-            response = requests.get(url)
-            init_image = Image.open(BytesIO(response.content)).convert("RGB")
-            init_image = init_image.resize((768, 512))
-     
         if use_vae:
             vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
             pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
@@ -92,9 +84,8 @@ def generate(
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
-        images = pipe(
+        return pipe(
             prompt=prompt,
-            image=init_image,
             negative_prompt=negative_prompt,
             prompt_2=prompt_2,
             negative_prompt_2=negative_prompt_2,
@@ -105,7 +96,7 @@ def generate(
             generator=generator,
             output_type="pil",
         ).images[0]
-        return images
+        return image
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
@@ -129,6 +120,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
                 step=0.01,
                 value=0.7,
             )
+        upload = gr.Image(label='Image2Image')
         with gr.Row():
             prompt = gr.Text(
                 placeholder="Input prompt",