35 lines
1.2 KiB
Python
35 lines
1.2 KiB
Python
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")
|