remove cuda

app.py

@@ -24,7 +24,7 @@ def generate(images, prompt, processor, model, device, dtype, generation_config)
 
 def main():
     # step 1: Setup constant
-    device = "cuda"
+    #device = "cuda"
     dtype = torch.float16
 
     # step 2: Load Processor and Model

app2.py

@@ -0,0 +1,48 @@
+import io
+import requests
+import torch
+from PIL import Image
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig
+
+def download_image(url):
+    resp = requests.get(url)
+    resp.raise_for_status()
+    return Image.open(io.BytesIO(resp.content)).convert("RGB")
+
+def generate(images, prompt, processor, model, device, dtype, generation_config):
+    inputs = processor(
+        images=images[:2], text=f" USER: <s>{prompt} ASSISTANT: <s>", return_tensors="pt"
+    ).to(device=device, dtype=dtype)
+    output = model.generate(**inputs, generation_config=generation_config)[0]
+    response = processor.tokenizer.decode(output, skip_special_tokens=True)
+    return response
+
+def main():
+    st.title("Medical Image Analysis")
+    
+    device = "cuda"
+    dtype = torch.float16
+
+    processor = AutoProcessor.from_pretrained("StanfordAIMI/CheXagent-8b", trust_remote_code=True)
+    generation_config = GenerationConfig.from_pretrained("StanfordAIMI/CheXagent-8b")
+    model = AutoModelForCausalLM.from_pretrained(
+        "StanfordAIMI/CheXagent-8b", torch_dtype=dtype, trust_remote_code=True
+    ).to(device)
+
+    image_path = "https://upload.wikimedia.org/wikipedia/commons/3/3b/Pleural_effusion-Metastatic_breast_carcinoma_Case_166_%285477628658%29.jpg"
+    images = [download_image(image_path)]
+
+    anatomies = [
+        "Airway", "Breathing", "Cardiac", "Diaphragm",
+        "Everything else (e.g., mediastinal contours, bones, soft tissues, tubes, valves, and pacemakers)"
+    ]
+
+    for anatomy in anatomies:
+        prompt = f'Describe "{anatomy}"'
+        response = generate(images, prompt, processor, model, device, dtype, generation_config)
+        st.subheader(f"Findings for [{anatomy}]:")
+        st.write(response)
+
+if __name__ == '__main__':
+    main()