Update app.py

app.py

@@ -9,8 +9,10 @@ import gradio as gr
 import numpy as np
 import PIL.Image
 import spaces
+import requests
 import torch
-from diffusers import AutoencoderKL, DiffusionPipeline
+from io import BytesIO
+from diffusers import StableDiffusionImg2ImgPipeline, AutoencoderKL, DiffusionPipeline
 
 DESCRIPTION = "# SDXL"
 if not torch.cuda.is_available():
@@ -55,12 +57,13 @@ def generate(
     model = 'SG161222/Realistic_Vision_V6.0_B1_noVAE',
     vaecall = 'stabilityai/sd-vae-ft-mse',
     lora = 'amazonaws-la/juliette',
+    url = 'https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg'
     lora_scale: float = 0.7,
 ) -> PIL.Image.Image:
     if torch.cuda.is_available():
 
         if not use_vae:  
-            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
+            pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model, torch_dtype=torch.float16)
             
         if use_vae:
             vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
@@ -69,7 +72,11 @@ def generate(
         if use_lora:
             pipe.load_lora_weights(lora)
             pipe.fuse_lora(lora_scale=0.7)
-                                   
+
+        response = requests.get(url)
+        init_image = Image.open(BytesIO(response.content)).convert("RGB")
+        init_image = init_image.resize((1024, 1024))
+
         if ENABLE_CPU_OFFLOAD:
             pipe.enable_model_cpu_offload()
             
@@ -99,6 +106,7 @@ def generate(
             guidance_scale=guidance_scale_base,
             num_inference_steps=num_inference_steps_base,
             generator=generator,
+            image=init_image
             output_type="pil",
         ).images[0]
     else: