Update app.py

app.py

@@ -8,6 +8,14 @@ import spaces
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+if torch.cuda.is_available():
+    device = "cuda"
+    print("Using GPU")
+else:
+    device = "cpu"
+    print("Using CPU")
+
+# download sd3 medium weights
 model_path = snapshot_download(
     repo_id="stabilityai/stable-diffusion-3-medium", 
     revision="refs/pr/26",
@@ -17,24 +25,19 @@ model_path = snapshot_download(
     token=HF_TOKEN,
     )
 
-if torch.cuda.is_available():
-    device = "cuda"
-    print("Using GPU")
-else:
-    device = "cpu"
-    print("Using CPU")
 
 # Initialize the pipeline and download the model
 pipe = StableDiffusion3Pipeline.from_pretrained(model_path, torch_dtype=torch.float16)
 pipe.to(device)
 
+# superprompt-v1
 tokenizer = T5Tokenizer.from_pretrained("roborovski/superprompt-v1")
 model = T5ForConditionalGeneration.from_pretrained("roborovski/superprompt-v1", device_map="auto", torch_dtype="auto")
 model.to(device)
 
 # Define the image generation function
 @spaces.GPU(duration=60)
-def generate_image(prompt, negative_prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt):
+def generate_image(prompt, enhance_prompt, negative_prompt, num_inference_steps, height, width, guidance_scale, seed, num_images_per_prompt):
     if seed == 0:
         seed = random.randint(1, 2**32-1)