updated code

app.py

@@ -1,21 +1,27 @@
-import gradio as gr
-from diffusers import StableDiffusionPipeline
 import torch
+from diffusers import StableDiffusionImg2ImgPipeline
+from PIL import Image
+import gradio as gr
 
-# Load the model (Make sure to set up the correct model path)
-pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
-# pipe = pipe.to("cuda")
-pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-
+# Load Stable Diffusion Image-to-Image Pipeline
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+    torch_dtype=torch.float16
+)
+pipe.to("cuda" if torch.cuda.is_available() else "cpu")  # Use GPU if available
 
-# Define the function to process the uploaded image
 def generate_headshot(image):
-    # Process the uploaded image and generate a professional headshot
-    generated_image = pipe(image).images[0]
-    return generated_image
+    # Convert to PIL format
+    image = Image.open(image).convert("RGB")
 
-# Create the Gradio interface
-iface = gr.Interface(fn=generate_headshot, inputs=gr.Image(), outputs=gr.Image(), title="Professional Headshot Generator")
+    # Define the text prompt to guide AI
+    prompt = "A professional corporate headshot, studio lighting, high resolution, DSLR quality"
+
+    # Run the pipeline with the uploaded image as input
+    generated_image = pipe(prompt=prompt, image=image, strength=0.7).images[0]
+
+    return generated_image
 
-# Launch the interface
+# Create Gradio UI
+iface = gr.Interface(fn=generate_headshot, inputs="image", outputs="image")
 iface.launch()