Update app.py

app.py

@@ -63,7 +63,14 @@ def generate(
         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 = StableDiffusionImg2ImgPipeline.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((1024, 1024))
+
         if use_lora:
             pipe.load_lora_weights(lora)
             pipe.fuse_lora(lora_scale)
@@ -86,6 +93,7 @@ def generate(
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
+    if not use_img2img:
         return pipe(
             prompt=prompt,
             negative_prompt=negative_prompt,
@@ -98,7 +106,21 @@ def generate(
             generator=generator,
             output_type="pil",
         ).images[0]
-        return image
+    else:
+        images = pipe(
+            prompt=prompt,
+            image=init_image,
+            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]
+        return images
 
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",