Update app.py

app.py

@@ -3,13 +3,41 @@ import torch
 import io
 import smtplib
 from email.message import EmailMessage
-from diffusers import StableDiffusionImg2ImgPipeline
+from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
 
-# Load SD3.5 pipeline
-pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-3.5-large", torch_dtype=torch.bfloat16
+from diffusers.utils import load_image
+from PIL import Image
+
+
+# Load the base (text-to-image) model
+pipe_t2i = StableDiffusionXLPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
 ).to("cuda")
 
+# Use from_pipe to initialize img2img from the same base
+pipe_i2i = StableDiffusionXLImg2ImgPipeline.from_pipe(pipe_t2i).to("cuda")
+
+def generate_from_image(image: Image.Image, prompt: str,
+                        strength: float = 0.75,
+                        guidance_scale: float = 10.5,
+                        steps: int = 40):
+    """
+    Generates a new SDXL image conditioned on both prompt and input image.
+    """
+    print(f"Prompt: {prompt}, Strength: {strength}, Guidance: {guidance_scale}")
+
+    image = image.convert("RGB").resize((1024, 1024))
+    with torch.inference_mode():
+        output = pipe_i2i(
+            prompt=prompt,
+            image=image,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=steps
+        ).images[0]
+    return output
+
 style_prompts = {
     "Single Photo With Dhaka Topi - Ghibli Style": "Create a Ghibli-style portrait using the uploaded photo as the person’s face to get a magical Ghibli-style portrait, wearing a traditional Nepali Dhaka topi. The dreamy background features Mount Everest, blooming laliguras flowers, Kathmandu’s Ghanta Ghar tower, elegant Newari window patterns, and colorful Nepali prayer flags waving in the wind. The art should blend Studio Ghibli’s enchanting style with authentic Nepali culture.",
     "Two Person Picture With Nepali Style Ghibli Art": "Create a Ghibli-style portrait using the uploaded photo as the person’s face. Background features Mount Everest, laliguras flowers, Kathmandu’s Ghanta Ghar tower, Newari window patterns, colorful Nepali prayer flags, the Nepali flag, Pashupatinath temple, and a flying Himalayan monal (Danfe). Make it vibrant, enchanting, and full of Nepali cultural details.",
@@ -20,18 +48,11 @@ style_prompts = {
 }
 sample_keys = list(style_prompts.keys())[:5]
 
-def generate_ghibli_style(image, prompt, strength=0.6):
-    with torch.inference_mode():
-        out = pipe(prompt=prompt,
-                   image=image,
-                   strength=strength,
-                   guidance_scale=6.5,
-                   num_inference_steps=25).images[0]
-    return out
+
 
 def process(image, style, email):
     prompt = style_prompts.get(style, "")
-    result = generate_ghibli_style(image, prompt)
+    result = generate_from_image(image, prompt)
     email_status = ""
     if email:
         email_status = send_email(email, result)