updated

app.py

@@ -28,39 +28,32 @@
 # iface = gr.Interface(fn=generate_headshot, inputs="image", outputs="image")
 # iface.launch()
 # ===================================
-# Install required packages (run in terminal if not installed)
-# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-# pip install diffusers transformers accelerate safetensors Pillow numpy
-
 import torch
-import numpy as np
-from PIL import Image
 from diffusers import StableDiffusionImg2ImgPipeline
+from PIL import Image
+import gradio as gr
 
-# Load Stable Diffusion 2.1 img2img model
-model_id = "stabilityai/stable-diffusion-2-1"
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id, torch_dtype=torch.float32)  
-# pipe.to("cpu")  # Change to "cuda" for GPU
-pipe.to("cuda" if torch.cuda.is_available() else "cpu")
-
-# Enable CPU optimizations
-pipe.enable_attention_slicing()
-pipe.enable_sequential_cpu_offload()
-
-# Load your sample image (make sure it's a clear, high-quality image)
-input_image = Image.open("your_sample_image.jpg").convert("RGB")
-
-# Resize image to match model's expected size (512x512)
-input_image = input_image.resize((512, 512))
+# Load Stable Diffusion Image-to-Image Pipeline
+pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-2-1",
+    torch_dtype=torch.float32  # Use float32 for CPU
+)
 
-# Define a denoising strength (0.2-0.8) → Lower = keeps original details, Higher = more AI changes
-denoising_strength = 0.5
+# Force execution on CPU (since no GPU is available)
+pipe.to("cpu")
 
-# Generate a professional headshot using img2img
-generated_image = pipe(image=input_image, strength=denoising_strength, guidance_scale=7.5).images[0]
+# Define function to generate professional headshots
+def generate_headshot(image):
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
 
-# Save and display the generated image
-generated_image.save("ai_headshot.png")
-generated_image.show()
+    prompt = "A professional corporate headshot, studio lighting, high resolution, DSLR quality"
+    
+    # Generate image (Lower strength for subtle changes)
+    generated_image = pipe(prompt=prompt, image=image, strength=0.5).images[0]
 
+    return generated_image
 
+# Create Gradio UI
+iface = gr.Interface(fn=generate_headshot, inputs="image", outputs="image")
+iface.launch()