import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image
import requests
from io import BytesIO

# Load the pipeline
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda" if torch.cuda.is_available() else "cpu")

# Function to load image from local path or URL
def load_image(path_or_url):
    if path_or_url.startswith("http"):
        response = requests.get(path_or_url)
        return Image.open(BytesIO(response.content)).convert("RGB")
    else:
        return Image.open(path_or_url).convert("RGB")

# Load original plant image
input_image = st.file_uploader("Upload an image of your plant", type=["jpg", "jpeg", "png"])

# Resize to 512x512 for best results
input_image = input_image.resize((512, 512))

# Define environmental condition prompt
prompt = "A withered plant with brown leaves, dry soil, and signs of dehydration, after 10 days of heat and no water"

# Generate the modified image
output = pipe(prompt=prompt, image=input_image, strength=0.75, guidance_scale=7.5)
output_image = output.images[0]

# Save the result
output_image.save("plant_future_prediction.png")