Update app.py

app.py

@@ -1,9 +1,10 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, StableDiffusionImg2ImgPipeline
 from transformers.utils.hub import move_cache
 import torch
+from PIL import Image
 
 move_cache()
 
@@ -11,17 +12,28 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 if torch.cuda.is_available():
     torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained("Envvi/Inkpunk-Diffusion", torch_dtype=torch.float16, variant="fp16")
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained("Envvi/Inkpunk-Diffusion", torch_dtype=torch.float16)
+    #pipe = DiffusionPipeline.from_pretrained("Envvi/Inkpunk-Diffusion", torch_dtype=torch.float16, variant="fp16")
     pipe.enable_xformers_memory_efficient_attention()
     pipe = pipe.to(device)
 else: 
-    pipe = DiffusionPipeline.from_pretrained("Envvi/Inkpunk-Diffusion")
+    pipe = StableDiffusionImg2ImgPipeline.from_pretrained("Envvi/Inkpunk-Diffusion", torch_dtype=torch.float16)
     pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def generate_image(uploaded_image):
+    # Open the uploaded image
+    image = Image.open(uploaded_image)
+    
+    # Run the image through the Stable Diffusion pipeline
+    with torch.no_grad():
+        output = pipe(image, guidance_scale=7.5)["sample"][0]
+    
+    return output
+
+def infer(base_img, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, strength):
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -29,10 +41,11 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
     generator = torch.Generator().manual_seed(seed)
     
     image = pipe(
+        image = base_img,
         prompt = prompt, 
         negative_prompt = negative_prompt,
         guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
+        strength = strength, 
         width = width, 
         height = height,
         generator = generator
@@ -65,6 +78,10 @@ with gr.Blocks(css=css) as demo:
         # Text-to-Image Gradio Template
         Currently running on {power_device}.
         """)
+
+
+        with gr.Row():
+            base_img = gr.Interface(fn=generate_image, inputs=gr.inputs.Image(type="pil"), outputs=gr.outputs.Image(type="pil"))            
         
         with gr.Row():
             
@@ -127,10 +144,10 @@ with gr.Blocks(css=css) as demo:
                     value=0.0,
                 )
                 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
+                strength = gr.Slider(
+                    label="strength",
                     minimum=1,
-                    maximum=12,
+                    maximum=10,
                     step=1,
                     value=2,
                 )
@@ -142,7 +159,7 @@ with gr.Blocks(css=css) as demo:
 
     run_button.click(
         fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps],
+        inputs = [base_img, prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, strength],
         outputs = [result]
     )