Update app.py

app.py

@@ -5,7 +5,7 @@ import numpy as np
 import spaces
 import random
 import string
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
 import torch
 from pathos.multiprocessing import ProcessingPool as ProcessPoolExecutor
 import requests
@@ -15,15 +15,15 @@ pool = ProcessPoolExecutor(64)
 pool.__enter__()
 
 #model_id = "runwayml/stable-diffusion-v1-5"
-model_id = "stabilityai/stable-diffusion-3-medium"
+model_id = "stabilityai/stable-diffusion-3-medium-diffusers"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 if torch.cuda.is_available():
     torch.cuda.max_memory_allocated(device=device)
-    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
+    pipe = StableDiffusion3Pipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
     pipe = pipe.to(device)
 else: 
-    pipe = DiffusionPipeline.from_pretrained(model_id, use_safetensors=True)
+    pipe = StableDiffusion3Pipeline.from_pretrained(model_id, use_safetensors=True)
     pipe = pipe.to(device)
 
 def translate(text,lang):
@@ -67,8 +67,8 @@ def infer(prompt):
         height=4096,
         width=4096,
         negative_prompt="",
-        num_inference_steps=50,
-        guidance_scale=9.0
+        num_inference_steps=100,
+        guidance_scale=10.0
     ).images[0].save(name)
     return name