Update app.py

app.py

@@ -48,6 +48,11 @@ def load_pipeline(model_name):
         StableDiffusionXLPipeline.from_single_file
         if MODEL.endswith(".safetensors")
         else StableDiffusionXLPipeline.from_pretrained
+    
+    img_pipeline = (
+        StableDiffusionXLImg2ImgPipeline.from_single_file
+        if MODEL.endswith(".safetensors")
+        else StableDiffusionXLImg2ImgPipeline.from_pretrained
     )
 
     pipe = pipeline(
@@ -60,9 +65,28 @@ def load_pipeline(model_name):
         use_auth_token=HF_TOKEN,
     )
 
+    img_pipe = img_pipeline(
+        model_name,
+        vae=vae,
+        torch_dtype=torch.float16,
+        custom_pipeline="lpw_stable_diffusion_xl",
+        use_safetensors=True,
+        add_watermarker=False,
+        use_auth_token=HF_TOKEN,
+    )
+
     pipe.to(device)
-    return pipe
+    img_pipe.to(device)
+    return pipe, img_pipe
 
+def load_img(resize_width,img):
+    img = Image.open(img)
+    width, height = img.size
+    scale = resize_width / width
+    resize_height = height * scale
+    
+    img = img.resize((resize_width, resize_height), Image.Resampling.LANCZOS)
+    return img, resize_width, resize_height
 
 @spaces.GPU
 def generate(
@@ -82,6 +106,7 @@ def generate(
     upscale_by: float = 1.5,
     add_quality_tags: bool = True,
     isImg2Img: bool = True,
+    img_path: str= ""
     
     progress=gr.Progress(track_tqdm=True),
 ):
@@ -107,6 +132,9 @@ def generate(
     backup_scheduler = pipe.scheduler
     pipe.scheduler = utils.get_scheduler(pipe.scheduler.config, sampler)
 
+    img_backup_scheduler = img_pipe.scheduler
+    img_pipe.scheduler = utils.get_scheduler(img_pipe.scheduler.config, sampler)
+
     if use_upscaler:
         upscaler_pipe = StableDiffusionXLImg2ImgPipeline(**pipe.components)
     metadata = {
@@ -142,38 +170,77 @@ def generate(
 
     try:
         if use_upscaler:
-            latents = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                output_type="latent",
-            ).images
-            upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
-            images = upscaler_pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                image=upscaled_latents,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                strength=upscaler_strength,
-                generator=generator,
-                output_type="pil",
-            ).images
+            if isImg2Img:
+                img, img_width, img_height = load_img(512, img_path)
+                latents = img_pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=img_width,
+                    height=img_height,
+                    image=img,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    output_type="latent",
+                ).images
+                upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
+                images = upscaler_pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    image=upscaled_latents,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    strength=upscaler_strength,
+                    generator=generator,
+                    output_type="pil",
+                ).images
+            else:
+                latents = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    output_type="latent",
+                ).images
+                upscaled_latents = utils.upscale(latents, "nearest-exact", upscale_by)
+                images = upscaler_pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    image=upscaled_latents,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    strength=upscaler_strength,
+                    generator=generator,
+                    output_type="pil",
+                ).images
         else:
-            images = pipe(
-                prompt=prompt,
-                negative_prompt=negative_prompt,
-                width=width,
-                height=height,
-                guidance_scale=guidance_scale,
-                num_inference_steps=num_inference_steps,
-                generator=generator,
-                output_type="pil",
-            ).images
+            if isImg2Img:
+                img, img_width, img_height = load_img(512, img_path)
+                images = pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=img_width,
+                    height=img_height,
+                    image=img,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    output_type="pil",
+                ).images
+            else:
+                images = img_pipe(
+                    prompt=prompt,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale,
+                    num_inference_steps=num_inference_steps,
+                    generator=generator,
+                    output_type="pil",
+                ).images
 
         if images:
             image_paths = [
@@ -192,14 +259,15 @@ def generate(
         if use_upscaler:
             del upscaler_pipe
         pipe.scheduler = backup_scheduler
+        img_pipe.scheduler = img_backup_scheduler
         utils.free_memory()
 
 
 if torch.cuda.is_available():
-    pipe = load_pipeline(MODEL)
+    pipe, img_pipe = load_pipeline(MODEL)
     logger.info("Loaded on Device!")
 else:
-    pipe = None
+    pipe, img_pipe = None, None
 
 styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in config.style_list}
 quality_prompt = {
@@ -241,7 +309,9 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
                     )
                     image = gr.Image(
                         label="Image Input",
-                        visible=False
+                        visible=False,
+                        source="upload", 
+                        type="filepath"
                     )
                     with gr.Accordion(label="Quality Tags", open=True):
                         add_quality_tags = gr.Checkbox(
@@ -402,6 +472,8 @@ with gr.Blocks(css="style.css", theme="NoCrypt/[email protected]") as demo:
             upscaler_strength,
             upscale_by,
             add_quality_tags,
+            isImg2Img,
+            img_path,
         ],
         outputs=[result, gr_metadata],
         api_name="run",