Update main.py

main.py

@@ -1,32 +1,71 @@
 import gradio as gr
 import torch
-from diffusers import StableDiffusers3Pipeline
+from diffusers import StableDiffusionPipeline
 
 def image_generation(prompt):
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    pipeline = StableDiffusers3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers",
-                                                        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-                                                        text_encoder_3 = None,
-                                                        tokenizer_3 = None)
-    # pipeline.to(device)
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-3-medium",
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    )
+    #pipeline.to(device)
     pipeline.enable_model_cpu_offload()
-
+    
     image = pipeline(
-        prompt = prompt,
-        negative_prompt = "blurred, ugly, watermark, low resolution, blurry, nude",
-        num_inference_steps = 40,
+        prompt=prompt,
+        negative_prompt="blurred, ugly, watermark, low resolution, blurry, nude",
+        num_inference_steps=40,
         height=1024,
         width=1024,
         guidance_scale=8.0
     ).images[0]
 
-    image.show()
+    return image
 
-interface= gr.interface(
+interface = gr.Interface(
     fn=image_generation,
-    inputs = gr.Textbox(lines="2", placeholder="Enter Your Prompt ..."),
-    outputs = gr.Image(type="pil"),
-    title = "AI Text Generation By SD-3M"
+    inputs=gr.Textbox(lines=2, placeholder="Enter Your Prompt ..."),
+    outputs=gr.Image(type="pil"),
+    title="AI Text Generation By SD-3M"
 )
 
-interface.launch()
+interface.launch()
+
+
+
+
+
+
+
+# import gradio as gr
+# import torch
+# from diffusers import StableDiffusers3Pipeline
+
+# def image_generation(prompt):
+#     device = "cuda" if torch.cuda.is_available() else "cpu"
+#     pipeline = StableDiffusers3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers",
+#                                                         torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+#                                                         text_encoder_3 = None,
+#                                                         tokenizer_3 = None)
+#     # pipeline.to(device)
+#     pipeline.enable_model_cpu_offload()
+
+#     image = pipeline(
+#         prompt = prompt,
+#         negative_prompt = "blurred, ugly, watermark, low resolution, blurry, nude",
+#         num_inference_steps = 40,
+#         height=1024,
+#         width=1024,
+#         guidance_scale=8.0
+#     ).images[0]
+
+#     image.show()
+
+# interface= gr.interface(
+#     fn=image_generation,
+#     inputs = gr.Textbox(lines="2", placeholder="Enter Your Prompt ..."),
+#     outputs = gr.Image(type="pil"),
+#     title = "AI Text Generation By SD-3M"
+# )
+
+# interface.launch()