Update app.py

app.py

@@ -5,7 +5,7 @@ import spaces
 import torch
 from diffusers import DiffusionPipeline
 
-dtype = torch.bfloat16
+dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 pipe = DiffusionPipeline.from_pretrained("shuttleai/shuttle-jaguar", torch_dtype=dtype).to(device)
@@ -20,14 +20,14 @@ def infer(prompt, seed=42, randomize_seed=False, width=1024, height=1024, num_in
     """
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
     image = pipe(
             prompt = prompt, 
             width = width,
             height = height,
             num_inference_steps = num_inference_steps, 
             generator = generator,
-            guidance_scale=0.0
+            guidance_scale=2.0
     ).images[0] 
     return image, seed
  
@@ -72,7 +72,7 @@ Shuttle Jaguar is a text-to-image AI model designed to generate highly aesthetic
                 minimum=0,
                 maximum=MAX_SEED,
                 step=1,
-                value=0,
+                value=42,
             )
             
             randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
@@ -94,6 +94,9 @@ Shuttle Jaguar is a text-to-image AI model designed to generate highly aesthetic
                     step=32,
                     value=1024,
                 )
+
+                width  = min(width, MAX_IMAGE_SIZE)
+                height = min(height, MAX_IMAGE_SIZE)
             
             with gr.Row():