updated

app.py

@@ -1,29 +1,65 @@
+# import torch
+# import numpy as np
+# from diffusers import StableDiffusionImg2ImgPipeline
+# from PIL import Image
+# import gradio as gr
+
+# # 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
+
+# def generate_headshot(image):
+#     # Convert NumPy array to PIL Image
+#     if isinstance(image, np.ndarray):
+#         image = Image.fromarray(image)
+
+#     # Define the AI prompt for professional headshots
+#     prompt = "A professional corporate headshot, studio lighting, high resolution, DSLR quality"
+
+#     # Generate the AI-enhanced headshot
+#     generated_image = pipe(prompt=prompt, image=image, strength=0.7).images[0]
+
+#     return generated_image
+
+# # Create Gradio UI
+# 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 diffusers import StableDiffusionImg2ImgPipeline
 from PIL import Image
-import gradio as gr
+from diffusers import StableDiffusionImg2ImgPipeline
+
+# 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
+
+# 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")
 
-# 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
+# Resize image to match model's expected size (512x512)
+input_image = input_image.resize((512, 512))
 
-def generate_headshot(image):
-    # Convert NumPy array to PIL Image
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
+# Define a denoising strength (0.2-0.8) → Lower = keeps original details, Higher = more AI changes
+denoising_strength = 0.5
 
-    # Define the AI prompt for professional headshots
-    prompt = "A professional corporate headshot, studio lighting, high resolution, DSLR quality"
+# Generate a professional headshot using img2img
+generated_image = pipe(image=input_image, strength=denoising_strength, guidance_scale=7.5).images[0]
 
-    # Generate the AI-enhanced headshot
-    generated_image = pipe(prompt=prompt, image=image, strength=0.7).images[0]
+# Save and display the generated image
+generated_image.save("ai_headshot.png")
+generated_image.show()
 
-    return generated_image
 
-# Create Gradio UI
-iface = gr.Interface(fn=generate_headshot, inputs="image", outputs="image")
-iface.launch()