app.py

import gradio as gr
import torch
import asyncio
from transformers import pipeline
from diffusers import StableDiffusionPipeline

# Check if GPU is available
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load models with optimized settings
image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224", device=0 if torch.cuda.is_available() else -1)
text_generator = pipeline("text-generation", model="google/flan-t5-small", device=0 if torch.cuda.is_available() else -1)  # Faster model
image_generator = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base", torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
image_generator.to(device)  # Move to GPU if available

async def process_image(image):
    try:
        # Run classification asynchronously
        classification_task = asyncio.to_thread(image_classifier, image)
        classification = await classification_task
        object_name = classification[0]["label"]

        # Generate reuse idea asynchronously
        prompt = f"Suggest a creative way to reuse a {object_name} in a useful manner."
        text_task = asyncio.to_thread(text_generator, prompt, max_length=50)
        idea = (await text_task)[0]["generated_text"]

        # Generate image asynchronously
        image_task = asyncio.to_thread(image_generator, prompt)
        generated_image = (await image_task).images[0]

        return object_name, idea, generated_image

    except Exception as e:
        return "Error", f"An error occurred: {str(e)}", None

# Create Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),
    outputs=["text", "text", "image"],
    title="AI-Powered Waste Reuse",
    description="Upload an image of a waste item, and AI will suggest reuse ideas along with an AI-generated image."
)

iface.launch()