Update src/app.py

src/app.py

@@ -1,4 +1,3 @@
-# Code developed by: Ruslan Magana Vsevolodovna
 """Template Demo for IBM Granite Hugging Face spaces."""
 
 from collections.abc import Iterator
@@ -60,15 +59,24 @@ text_tokenizer.use_default_system_prompt = False
 # Vision model loading
 vision_model_path = "ibm-granite/granite-vision-3.1-2b-preview"
 vision_processor = LlavaNextProcessor.from_pretrained(vision_model_path, use_fast=True)
+
+# Option 1: Use the default settings (like the original demo)
 vision_model = LlavaNextForConditionalGeneration.from_pretrained(
     vision_model_path,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    low_cpu_mem_usage=True,
-    force_download=True, # Added force_download to ensure fresh download
-    revision="main" # Added revision to specify main branch (if needed)
+    torch_dtype="auto",
+    device_map="auto"
 )
 
+# --- OR ---
+
+# Option 2: Use torch.float16 but ensure the custom model code is trusted
+# vision_model = LlavaNextForConditionalGeneration.from_pretrained(
+#     vision_model_path,
+#     torch_dtype=torch.float16,
+#     device_map="auto",
+#     trust_remote_code=True
+# )
+
 
 @spaces.GPU
 def generate(
@@ -323,7 +331,7 @@ with gr.Blocks(fill_height=True, css_paths=css_file_path, head_paths=head_file_p
     )
 
 
-    image_examples_dir = Path(__file__).parent / "image_examples" # Create a folder 'image_examples' in the same directory as app.py and put images there
+    image_examples_dir = Path(__file__).parent / "themes" # Create a folder 'themes' in the same directory as app.py and put images there
     image_examples = [
         str(image_examples_dir / "cat.jpg"), # Replace "cat.jpg" with actual image file names in 'image_examples'
         str(image_examples_dir / "dog.jpg"), # Replace "dog.jpg" with actual image file names in 'image_examples'