Upload 4 files

app.py

@@ -10,7 +10,13 @@ from diffusers import StableDiffusionXLPipeline, EDMEulerScheduler, StableDiffus
 from custom_pipeline import CosStableDiffusionXLInstructPix2PixPipeline
 from huggingface_hub import hf_hub_download
 from huggingface_hub import InferenceClient
+from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel, FlowMatchEulerDiscreteScheduler
 
+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:
@@ -37,47 +43,19 @@ def set_timesteps_patched(self, num_inference_steps: int, device = None):
 
 # Image Editor
 edit_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl_edit.safetensors")
-normal_file = hf_hub_download(repo_id="stabilityai/cosxl", filename="cosxl.safetensors")
-
 EDMEulerScheduler.set_timesteps = set_timesteps_patched
-
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
-
 pipe_edit = StableDiffusionXLInstructPix2PixPipeline.from_single_file(
     edit_file, num_in_channels=8, is_cosxl_edit=True, vae=vae, torch_dtype=torch.float16,
 )
 pipe_edit.scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=120.0, sigma_data=1.0, prediction_type="v_prediction")
 pipe_edit.to("cuda")
 
-from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
-
-if not torch.cuda.is_available():
-    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo may not work on CPU.</p>"
-
-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-v4",
-        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",
+def king(type ,
+        input_image ,
+        instruction: str ,
         steps: int = 8,
         randomize_seed: bool = False,
         seed: int = 25,
@@ -85,12 +63,13 @@ def king(type = "Image Generation",
         image_cfg_scale: float = 1.7,
         width: int = 1024,
         height: int = 1024,
-        guidance_scale: float = 6.2,
+        guidance_scale: float = 6,
         use_resolution_binning: bool = True,
         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,49 +82,34 @@ 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 = instruction,
+            guidance_scale = 7, 
+            num_inference_steps = steps, 
+            width = width, 
+            height = height,
+            generator = generator
+        ).images[0]         
+        return seed, image
+
+client = InferenceClient()
 # Prompt classifier
-def response(instruction, input_image=None):
+def response(instruction, input_image=None ):
     if input_image is None:
         output="Image Generation"
-        yield output
     else:
-        client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-            
-        generate_kwargs = dict(
-                max_new_tokens=5,
-            )
-    
-        system="[SYSTEM] You will be provided with text, and your task is to classify task is image generation or image editing answer with only task do not say anything else and stop as soon as possible. [TEXT]"
-            
-        formatted_prompt = system + instruction + "[TASK]"
-        stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-        output = ""
-        for response in stream:        
-            if not response.token.text == "</s>":
-                output += response.token.text
-        if "editing" in output:
+        text = instruction
+        labels = ["Image Editing", "Image Generation"]
+        classification = client.zero_shot_classification(text, labels, multi_label=True)
+        output = classification[0]
+        output = str(output)
+        if "Editing" in output:
             output = "Image Editing"
         else:
             output = "Image Generation"
-        yield output
     return output
 
 css = '''
@@ -160,7 +124,7 @@ examples=[
         [
             "Image Generation",
             None,
-            "A Super Car",
+            "A luxurious supercar with a unique design. The car should have a pearl white finish, and gold accents. 4k, realistic.",
 
         ],
         [
@@ -178,7 +142,7 @@ examples=[
         [
             "Image Generation",
             None,
-            "Kids going o school, Anime style",
+            "Ironman fighting with hulk, wall painting",
 
         ],
         [

requirements.txt

@@ -5,6 +5,7 @@ numpy
 transformers
 accelerate
 safetensors
-diffusers
+git+https://github.com/huggingface/diffusers.git
 spaces
-peft
+peft
+sentencepiece