Update app.py

app.py

@@ -1,21 +1,34 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+
+# Initialize the image classification pipeline
+classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
+
+# Initialize the tokenizer and model for the generative text (GPT-like model)
+tokenizer = AutoTokenizer.from_pretrained("gpt2")  # Example model, replace with your choice
+model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+def generate_tweet(label):
+    # Generate a promotional tweet using a GPT-like model
+    prompt = f"Write a creative and promotional tweet about {label}:"
+    inputs = tokenizer.encode(prompt, return_tensors="pt")
+    outputs = model.generate(inputs, max_length=280, num_return_sequences=1)
+    tweet = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return tweet
 
 def predict(image):
-    model_id = "google/vit-base-patch16-224"
-    classifier = pipeline("image-classification", model=model_id)
     predictions = classifier(image)
     # Sort predictions based on confidence and select the top one
     top_prediction = sorted(predictions, key=lambda x: x['score'], reverse=True)[0]
+    label = top_prediction['label'].split(',')[0]  # Clean up label if necessary
     
-    # Generate a promotional tweet based on the top prediction
-    tweet_template = "Check out this amazing {label}! 📸✨ Explore more about it and let your curiosity lead you to discover wonders."
-    tweet_text = tweet_template.format(label=top_prediction['label'].split(',')[0])  # Using split to clean up label if necessary
-    return tweet_text
+    # Generate the tweet
+    tweet = generate_tweet(label)
+    return tweet
 
-title = "Image Classifier to Promotional Tweet"
-description = "This demo recognizes and classifies images using the 'google/vit-base-patch16-224' model. Below, you'll see a generated promotional tweet based on the top prediction. Your task: Upload an image, and let's write a tweet about it!"
+title = "Image Classifier to Generative Tweet"
+description = "This demo recognizes and classifies images using the 'google/vit-base-patch16-224' model and generates a creative promotional tweet about the top prediction using a generative text model."
 input_component = gr.Image(type="pil", label="Upload an image here")
-output_component = gr.Textbox(label="Generated Promotional Tweet", placeholder="Write a tweet about the image")
+output_component = gr.Textbox(label="Generated Promotional Tweet")
 
 gr.Interface(fn=predict, inputs=input_component, outputs=output_component, title=title, description=description).launch()