Update app.py

app.py

@@ -10,7 +10,18 @@ from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffus
 from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download
 from huggingface_hub import InferenceClient
+import gradio as gr
+import numpy as np
+import random
+import torch
+from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
+import spaces
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.float16
 
+repo = "stabilityai/stable-diffusion-3-medium-diffusers"
+pipe = StableDiffusion3Pipeline.from_pretrained(repo, torch_dtype=torch.float16).to(device)
 
 help_text = """
 To optimize image results:
@@ -56,28 +67,11 @@ if not torch.cuda.is_available():
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-
-# Image Generator
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(
-        "fluently/Fluently-XL-Final",
-        torch_dtype=torch.float16,
-        use_safetensors=True,
-    )
-    pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    pipe.load_lora_weights("ehristoforu/dalle-3-xl-v2", weight_name="dalle-3-xl-lora-v2.safetensors", adapter_name="dalle")
-    pipe.set_adapters("dalle")
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, 999999)
-    return seed
-
 # Generator
 @spaces.GPU(duration=30, queue=False)
 def king(type = "Image Generation",
         input_image = None,
-        instruction: str = "Eiffel tower",
+        instruction: str ,
         steps: int = 8,
         randomize_seed: bool = False,
         seed: int = 25,
@@ -90,7 +84,8 @@ def king(type = "Image Generation",
         progress=gr.Progress(track_tqdm=True),
     ):
     if type=="Image Editing" :
-        seed = int(randomize_seed_fn(seed, randomize_seed))
+        if randomize_seed:
+            seed = random.randint(0, 99999)
         text_cfg_scale = text_cfg_scale
         image_cfg_scale = image_cfg_scale
         input_image = input_image
@@ -103,23 +98,18 @@ def king(type = "Image Generation",
             num_inference_steps=steps, generator=generator).images[0]
         return seed, output_image
     else :
-        pipe.to(device)
-        seed = int(randomize_seed_fn(seed, randomize_seed))
-        generator = torch.Generator().manual_seed(seed)   
-        
-        options = {
-            "prompt":instruction,
-            "width":width,
-            "height":height,
-            "guidance_scale":guidance_scale,
-            "num_inference_steps":steps,
-            "generator":generator,
-            "use_resolution_binning":use_resolution_binning,
-            "output_type":"pil",
-        }
-        
-        output_image = pipe(**options).images[0]
-        return seed, output_image
+        if randomize_seed:
+            seed = random.randint(0, 99999)
+        generator = torch.Generator().manual_seed(seed)
+        image = pipe(
+            prompt = prompt,
+            guidance_scale = guidance_scale, 
+            num_inference_steps = steps, 
+            width = width, 
+            height = height,
+            generator = generator
+        ).images[0]         
+        return seed, image
 
 # Prompt classifier
 def response(instruction, input_image=None):