Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import Qwen2_5_VLForConditionalGeneration, AutoProcessor
 from qwen_vl_utils import process_vision_info
 from PIL import Image
 import torch
+import os
 
 # Load the model and processor
 model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
@@ -13,23 +14,21 @@ model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
 )
 processor = AutoProcessor.from_pretrained("daniel3303/QwenStoryteller")
 
-def upload_file(files):
-    file_paths = [file.name for file in files]
-    return file_paths
-    
 @spaces.GPU()
-def generate_story(images):
+@torch.no_grad()
+def generate_story(file_paths):
+    # Load images from the file paths
+    images = [Image.open(file_path) for file_path in file_paths]
+    
     image_content = []
-    for img in images[:6]:
+    for img in images[:6]:  # Limit to 6 images
         image_content.append({
             "type": "image",
             "image": img,
         })
 
-    # Add text prompt at the end
     image_content.append({"type": "text", "text": "Generate a story based on these images."})
 
-    # Create messages with system prompt
     messages = [
         {
             "role": "system",
@@ -41,7 +40,6 @@ def generate_story(images):
         }
     ]
 
-    # Preparation for inference
     text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     image_inputs, video_inputs = process_vision_info(messages)
     inputs = processor(
@@ -53,7 +51,6 @@ def generate_story(images):
     )
     inputs = inputs.to(model.device)
 
-    # Inference: Generate the output
     generated_ids = model.generate(
         **inputs,
         max_new_tokens=4096,
@@ -61,6 +58,7 @@ def generate_story(images):
         temperature=0.7,
         top_p=0.9
     )
+
     generated_ids_trimmed = [
         out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
     ]
@@ -74,18 +72,19 @@ def generate_story(images):
 
 with gr.Blocks() as demo:
     gr.Markdown("# Qwen Storyteller \n## Upload up to 6 images to generate a creative story.")
+    
     with gr.Row():
         with gr.Column():
-            file_output = gr.File()
-            upload_button = gr.UploadButton("Upload up to 6 images", file_types=["image", "video"], file_count="multiple")
-            gen_button = gr.Button("Generate", variant="secondary")
+            upload_button = gr.UploadButton("Upload up to 6 images", file_types=["image"], file_count="multiple")
+            output_file = gr.File(label="Uploaded Files")
 
         with gr.Column():
-            outputs=gr.Textbox(label="Generated Story", lines=10)
+            outputs = gr.Textbox(label="Generated Story", lines=10)
 
-    upload_button.upload(upload_file, upload_button, file_output)
-    gen_button.click(fn=generate_story, file_output, outputs)
-    
+    upload_button.upload(lambda files: [f.name for f in files], upload_button, output_file)
+
+    gen_button = gr.Button("Generate", variant="secondary")
+    gen_button.click(generate_story, upload_button, outputs)
 
 if __name__ == "__main__":
     demo.launch()