import gradio as gr
from PIL import Image
import torch
from diffusers import StableDiffusionImg2ImgPipeline

MODEL_ID = "runwayml/stable-diffusion-v1-5" 

device = "cuda" if torch.cuda.is_available() else "cpu"

try:
    print(f"Loading model on device: {device}")
    
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        MODEL_ID, 
        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
        low_cpu_mem_usage=True, # Added optimization for low RAM environments
    )
    pipe = pipe.to(device)
    print("Model loaded successfully.")

except Exception as e:
    print(f"Error loading model: {e}")
    pipe = None 
    print("Image generation disabled due to model loading error.")


def generate_avatar(
    initial_image: Image, 
    prompt: str, 
    style_strength: float, 
    steps: int,
    negative_prompt: str
):
    """
    Generates a new image based on an input image and a text prompt.
    The 'style_strength' controls how much the new image deviates from the original.
    """
    if pipe is None:
           if not prompt:
        return initial_image, "Please enter a descriptive text prompt."
        
    if initial_image is None:
        return initial_image, "Please upload a reference image."

    print(f"Generating image with prompt: '{prompt}' and strength: {style_strength}")
    
    try:
              with torch.autocast(device) if device == "cuda" else torch.no_grad():
            generated_image = pipe(
                prompt=prompt,
                image=initial_image.convert("RGB"), # Ensure image is in RGB format
                strength=style_strength,
                num_inference_steps=steps,
                guidance_scale=7.5,
                negative_prompt=negative_prompt
            ).images[0]
            
        return generated_image, "Avatar successfully generated!"
        
    except Exception as e:
        print(f"Inference error: {e}")
        return initial_image, f"An error occurred during generation: {e}"


# Define the input components
image_input = gr.Image(type="pil", label="1. Upload Base Image (Your Photo or Style Reference)")
prompt_input = gr.Textbox(
    label="2. Creative Prompt (e.g., 'A professional LinkedIn headshot, cyberpunk aesthetic, highly detailed')", 
    placeholder="Describe the desired style and subject."
)
strength_slider = gr.Slider(
    minimum=0.1, 
    maximum=0.9, 
    value=0.65, 
    step=0.05, 
    label="3. Transformation Strength (Higher = More stylistic change)"
)
steps_slider = gr.Slider(
    minimum=10, 
    maximum=100, 
    value=30, 
    step=5, 
    label="4. Inference Steps (Higher = Better quality, Slower generation)"
)
negative_prompt_input = gr.Textbox(
    label="5. Negative Prompt (What to AVOID, e.g., 'ugly, blurry, low quality, cropped')", 
    placeholder="Enter elements to avoid in the generated image."
)
gr.Interface(
    fn=generate_avatar,
    inputs=[image_input, prompt_input, strength_slider, steps_slider, negative_prompt_input],
    outputs=[gr.Image(type="pil", label="Generated Profile Picture"), gr.Textbox(label="Status")],
    title="✨ Creative Profile Picture Generator (SFW)",
    description="Upload a reference photo and a creative prompt to generate a new, stylized avatar using image-to-image diffusion. Transformation strength controls how much the output deviates from the input image (0.1 is subtle, 0.9 is highly stylized).",
    live=False,
    allow_flagging="never"
).launch()