Create app.py

app.py

@@ -0,0 +1,90 @@
+import os
+import time
+import requests
+from PIL import Image
+from io import BytesIO
+from selenium import webdriver
+from selenium.webdriver.chrome.service import Service
+from webdriver_manager.chrome import ChromeDriverManager
+from diffusers import StableDiffusionPipeline
+import torch
+import gradio as gr
+
+# ---------- Step 1: Scrape Celebrity Images ----------
+def scrape_images(celebrity_name, num_images=20):
+    search_url = f"https://www.google.com/search?q={celebrity_name}+portrait&tbm=isch"
+    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()))
+    driver.get(search_url)
+    os.makedirs(f"data/{celebrity_name}", exist_ok=True)
+
+    images = driver.find_elements("tag name", "img")
+    count = 0
+
+    for img in images:
+        if count >= num_images:
+            break
+        src = img.get_attribute("src")
+        if src and "http" in src:
+            try:
+                img_data = requests.get(src).content
+                with open(f"data/{celebrity_name}/{celebrity_name}_{count}.jpg", "wb") as handler:
+                    handler.write(img_data)
+                count += 1
+            except Exception as e:
+                print(f"Error downloading image: {e}")
+    driver.quit()
+
+# ---------- Step 2: Fine-Tuning Stable Diffusion ----------
+def fine_tune_sd3(celebrity_name):
+    model_id = "runwayml/stable-diffusion-v1-5"
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)
+
+    celeb_images_path = f"data/{celebrity_name}"
+    images = [Image.open(os.path.join(celeb_images_path, img)) for img in os.listdir(celeb_images_path) if img.endswith(".jpg")]
+
+    # Simple fine-tuning logic (for demonstration; deep fine-tuning requires more work)
+    print(f"Fine-tuning with {len(images)} images of {celebrity_name}...")
+
+    # Saving model
+    fine_tuned_model_path = f"models/{celebrity_name}_sd3"
+    os.makedirs(fine_tuned_model_path, exist_ok=True)
+    pipe.save_pretrained(fine_tuned_model_path)
+    print(f"Model saved at {fine_tuned_model_path}")
+
+    return fine_tuned_model_path
+
+# ---------- Step 3: Generate Phone Cover Designs ----------
+def generate_cover(prompt, model_path):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    pipe = StableDiffusionPipeline.from_pretrained(model_path).to(device)
+
+    image = pipe(prompt).images[0]
+    cover_template = Image.open("phone_cover_template.png").convert("RGBA")
+    image = image.resize(cover_template.size)
+    blended = Image.alpha_composite(cover_template, image.convert("RGBA"))
+    
+    output_path = "generated_phone_cover.png"
+    blended.save(output_path)
+    return output_path
+
+# ---------- Step 4: Gradio Deployment ----------
+def launch_gradio(model_path):
+    def infer(prompt):
+        result = generate_cover(prompt, model_path)
+        return result
+
+    gr.Interface(fn=infer, 
+                 inputs=gr.Textbox(label="Enter a design prompt"),
+                 outputs=gr.Image(label="Generated Phone Cover"),
+                 title="Celebrity Phone Cover Generator").launch()
+
+# ---------- Main Workflow ----------
+if __name__ == "__main__":
+    celebrity = "Taylor Swift"  # Example celebrity
+    scrape_images(celebrity, num_images=30)
+    
+    model_path = fine_tune_sd3(celebrity)
+
+    # Deploy on Gradio
+    launch_gradio(model_path)