Update app.py

app.py

@@ -7,22 +7,19 @@ from diffusers import DiffusionPipeline
 import torch
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
 
 if torch.cuda.is_available():
-    torch_dtype = torch.float16
+    torch_dtype = torch.bfloat16
 else:
     torch_dtype = torch.float32
 
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 
 # @spaces.GPU #[uncomment to use ZeroGPU]
 def infer(
+    model_repo_id,
     prompt,
     negative_prompt,
     seed,
@@ -38,6 +35,9 @@ def infer(
 
     generator = torch.Generator().manual_seed(seed)
 
+    
+    pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+    pipe = pipe.to(device)
     image = pipe(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -52,9 +52,9 @@ def infer(
 
 
 examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
+    "mimic: PA view, FINAL REPORT\n PA AND LATERAL CHEST OF ___\n \n COMPARISON:  ___ radiograph.\n \n FINDINGS:  Cardiac silhouette is upper limits of normal in size and\n accompanied by pulmonary vascular congestion and a basilar predominant\n interstitial abnormality which most likely represents interstitial edema. \n Small bilateral pleural effusions are present, left greater than right, with\n interval decrease in size since the prior radiograph.  There is also improving\n aeration in the left retrocardiac region, likely resolving atelectasis.",
+    "mimic: AP view, small left-sided pleural effusion, discrete density overlying the postero-lateral rib, supraclavicular catheter at the cavoatrial junction, new opacity within the left lower lobe and lingula",
+    "mimic: AP view, FINAL REPORT\n EXAMINATION:  CHEST (PORTABLE AP)\n \n INDICATION:  ___ year old woman with ?TIA  // r/o acute CP process      r/o\n acute CP process\n \n IMPRESSION: \n \n In comparison with the study of ___, there again are low lung\n volumes with elevation of the right hemidiaphragmatic contour. Cardiac\n silhouette is at the upper limits of normal or mildly enlarged. No evidence of\n acute pneumonia or vascular congestion.\n \n Right IJ catheter extends to the mid to lower portion of the SVC.",
 ]
 
 css = """
@@ -66,7 +66,14 @@ css = """
 
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
+        gr.Markdown(" # SD Demo")
+
+        with gr.Row():
+            model_repo_id = gr.Dropdown(
+                choices=['Cylumn/sd2-findings', 'Cylumn/sd2-report', 'Cylumn/sd3-report'],
+                label="Select Model",
+                value='Cylumn/sd3-report'
+            )
 
         with gr.Row():
             prompt = gr.Text(
@@ -102,7 +109,7 @@ with gr.Blocks(css=css) as demo:
             with gr.Row():
                 width = gr.Slider(
                     label="Width",
-                    minimum=256,
+                    minimum=512,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
                     value=1024,  # Replace with defaults that work for your model
@@ -110,7 +117,7 @@ with gr.Blocks(css=css) as demo:
 
                 height = gr.Slider(
                     label="Height",
-                    minimum=256,
+                    minimum=512,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
                     value=1024,  # Replace with defaults that work for your model
@@ -119,10 +126,10 @@ with gr.Blocks(css=css) as demo:
             with gr.Row():
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
-                    minimum=0.0,
+                    minimum=7.5,
                     maximum=10.0,
                     step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
+                    value=7.5,  # Replace with defaults that work for your model
                 )
 
                 num_inference_steps = gr.Slider(
@@ -130,7 +137,7 @@ with gr.Blocks(css=css) as demo:
                     minimum=1,
                     maximum=50,
                     step=1,
-                    value=2,  # Replace with defaults that work for your model
+                    value=35,  # Replace with defaults that work for your model
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])
@@ -138,6 +145,7 @@ with gr.Blocks(css=css) as demo:
         triggers=[run_button.click, prompt.submit],
         fn=infer,
         inputs=[
+            model_repo_id,
             prompt,
             negative_prompt,
             seed,