Update app.py

app.py

@@ -1,6 +1,5 @@
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, VisionEncoderDecoderModel, ViTImageProcessor, pipeline
+from transformers import AutoTokenizer, VisionEncoderDecoderModel, ViTImageProcessor
 from diffusers import StableDiffusionPipeline
 
 # Initialize device and models for captioning
@@ -16,32 +15,38 @@ caption_model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(d
 diffusion_model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
 diffusion_model = diffusion_model.to(device)
 
-def predict(image):
+def get_caption(image):
     # Generate a caption from the image
     image = image.convert('RGB')
     image_tensor = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
     caption_ids = caption_model.generate(image_tensor, max_length=128, num_beams=3)[0]
     caption_text = tokenizer.decode(caption_ids, skip_special_tokens=True)
     return caption_text
-    
-def generate(image):
-    caption=predict(image)
+
+def generate_image(caption):
     # Generate an image from the caption
     generated_image = diffusion_model(caption)["sample"][0]
+    return generated_image
 
-    return caption, generated_images[0]
-    
 # Set up Gradio interface
-input = gr.Image(label="Upload any Image", type='pil')
-outputs = [gr.Textbox(label="Caption"), gr.Image(label="Generated Image")]
-#examples = ['example1.jpeg']
+image_input = gr.Image(label="Upload any Image", type='pil')
+caption_output = gr.Textbox(label="Caption")
+generated_image_output = gr.Image(label="Generated Image")
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            image = image_input()
+            get_caption_btn = gr.Button("Get Caption")
+        with gr.Column():
+            caption = caption_output()
+            generate_image_btn = gr.Button("Generate Image")
+    with gr.Row():
+        generated_image = generated_image_output()
+
+    caption.update(get_caption_btn.click(get_caption, inputs=image, outputs=caption))
+    generated_image.update(generate_image_btn.click(generate_image, inputs=caption, outputs=generated_image))
 
 title = "Image Captioning and Generation"
-interface = gr.Interface(
-    fn=generate,
-    inputs=input,
-    outputs=outputs,
-    #examples=examples,
-    title=title,
-)
-interface.launch()
+demo.launch(title=title)
+